adam_update_tritonupdate_fn_kernel_correctnessMathematical execution | Original-kernel candidate Masked per-program Lion parameter and momentum update (the Python name says Adam). Separate parameter/momentum regions; no host/grid composition or IEEE claim. | Statement and preconditionsspecification update_fn_kernel_correctness
(p_ptr grad_ptr exp_avg_ptr : RegionName)
(lr wd beta1 beta2 : ℝ) (n_elements BLOCK_SIZE : Nat)
(hRegions : p_ptr ≠ exp_avg_ptr) :
adamIO p_ptr grad_ptr exp_avg_ptr lr wd beta1 beta2 n_elements BLOCK_SIZE
⊨ fun p grad expAvg =>
(fun i => TiledOptimizer.lionParam (p i) (expAvg i) (grad i) lr wd beta1,
fun i => TiledOptimizer.lionMomentum (expAvg i) (grad i) beta2)
IO definition: adamIOdef adamIO (p_ptr grad_ptr exp_avg_ptr : RegionName)
(lr wd beta1 beta2 : ℝ) (n_elements BLOCK_SIZE : Nat) :
MaskedKernelIO₃ₓ₂ where
kernel := update_fn_kernel p_ptr grad_ptr exp_avg_ptr
lr wd beta1 beta2 n_elements BLOCK_SIZE
bufs := [p_ptr, grad_ptr, exp_avg_ptr] -- p and exp_avg are updated in place
in1 := p_ptr
in2 := grad_ptr
in3 := exp_avg_ptr
out1 := p_ptr -- = in1: in-place parameter update
out2 := exp_avg_ptr -- = in3: in-place momentum update
B := BLOCK_SIZE
read1 := fun pid => pid * BLOCK_SIZE
read2 := fun pid => pid * BLOCK_SIZE
read3 := fun pid => pid * BLOCK_SIZE
write1 := fun pid => pid * BLOCK_SIZE
write2 := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < n_elements
Referenced model/specification definitions: Python source |
|---|
add_exampleadd_kernel_correctnessMathematical execution | Original-kernel candidate Masked per-program sum for arbitrary element count and block size; disjoint allocations and active-lane bounds are part of the IO contract. | Statement and preconditionsspecification add_kernel_correctness
(in_ptr0 in_ptr1 out_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) :
addIO in_ptr0 in_ptr1 out_ptr n_elements BLOCK_SIZE
⊨ fun xs ys i => xs i + ys i
IO definition: addIOdef addIO (in_ptr0 in_ptr1 out_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) : MaskedKernelIO₂ where
kernel := add_kernel in_ptr0 in_ptr1 out_ptr n_elements BLOCK_SIZE
in1 := in_ptr0
in2 := in_ptr1
out := out_ptr
B := BLOCK_SIZE
read1 := fun pid => pid * BLOCK_SIZE
read2 := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < n_elements
Referenced model/specification definitions: Python source |
|---|
add_valuepuzzle1_kernel_correctnessMathematical execution | Original-kernel candidate Masked per-program scalar addition. Separate exact and abstract store-rounding contracts; no concrete hardware rounding claim. | Statement and preconditionsspecification puzzle1_kernel_correctness
(x_ptr output_ptr : RegionName)
(N BLOCK_SIZE : Nat) (value : ℝ) :
addValueIO x_ptr output_ptr N BLOCK_SIZE value
⊨ fun xs i => xs i + value
IO definition: addValueIOdef addValueIO (x_ptr output_ptr : RegionName)
(N BLOCK_SIZE : Nat) (value : ℝ) : MaskedKernelIO₁ where
kernel := puzzle1_kernel x_ptr output_ptr N BLOCK_SIZE value
inp := x_ptr
out := output_ptr
B := BLOCK_SIZE
read := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < N
Referenced model/specification definitions: Python source |
|---|
add_valuepuzzle1_kernel_io_correctnessAbstract cast/store rounding | Original-kernel candidate Masked per-program scalar addition. Separate exact and abstract store-rounding contracts; no concrete hardware rounding claim. | Statement and preconditionsspecification puzzle1_kernel_io_correctness (R : RoundingModel)
(x_ptr output_ptr : RegionName)
(N BLOCK_SIZE : Nat) (value : ℝ) :
addValueIO x_ptr output_ptr N BLOCK_SIZE value
⊨[R, .real] fun xs i => xs i + value
IO definition: addValueIOdef addValueIO (x_ptr output_ptr : RegionName)
(N BLOCK_SIZE : Nat) (value : ℝ) : MaskedKernelIO₁ where
kernel := puzzle1_kernel x_ptr output_ptr N BLOCK_SIZE value
inp := x_ptr
out := output_ptr
B := BLOCK_SIZE
read := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < N
Referenced model/specification definitions: Python source |
|---|
apply_penaltyapply_penalty_correctnessMathematical execution | Original-kernel candidate Per-batch penalty gather/scatter using natural-valued token metadata. Bounds and non-colliding output windows come from the IO contract; host metadata construction is external. | Statement and preconditionsspecification apply_penalty_correctness
(Logits presence_penalty freqency_penalty repetition_penalty : Region .real)
(p_token_ids p_token_counts p_cumsum_seq_len : Region .nat)
(stride_logit_b stride_logit_s BLOCK_P : Nat) :
applyPenaltyIO Logits presence_penalty freqency_penalty
repetition_penalty p_token_ids p_token_counts p_cumsum_seq_len
stride_logit_b stride_logit_s BLOCK_P
⊨ fun _ _ g₁ g₂ g₃ _ _ _ cnts xs j =>
penaltyValuePure g₁ g₂ g₃ cnts xs j
IO definition: applyPenaltyIOdef applyPenaltyIO
(Logits presence_penalty freqency_penalty repetition_penalty : Region .real)
(p_token_ids p_token_counts p_cumsum_seq_len : Region .nat)
(stride_logit_b stride_logit_s BLOCK_P : Nat) :
MetaScatterMasked2DKernelIO₁ where
kernel := apply_penalty Logits presence_penalty freqency_penalty
repetition_penalty p_token_ids p_token_counts p_cumsum_seq_len
stride_logit_b stride_logit_s BLOCK_P
fbuf1 := presence_penalty
fbuf2 := freqency_penalty
fbuf3 := repetition_penalty
mbuf := p_cumsum_seq_len
idbuf := p_token_ids
cntbuf := p_token_counts
inp := Logits
out := Logits
B := BLOCK_P
fwin1 := fun pid₀ _ => pid₀
fwin2 := fun pid₀ _ => pid₀
fwin3 := fun pid₀ _ => pid₀
mwin1 := fun pid₀ _ => pid₀
mwin2 := fun pid₀ _ => pid₀ + 1
readi := fun _ _ m₁ _ j => m₁ + j.val
readc := fun _ _ m₁ _ j => m₁ + j.val
read := fun pid₀ _ _ _ ids j => pid₀ * stride_logit_b + ids j
mask := fun _ _ m₁ m₂ j => m₁ + j.val < m₂
write := fun pid₀ _ _ _ ids j => pid₀ * stride_logit_b + ids j
Referenced model/specification definitions: Python source |
|---|
attention_forward_tritonattention_forward_triton_closed_form_correctMathematical execution | Configured model / stage Forward attention with the explicitly shared contiguous head strides and a whole number of KV blocks. The IO theorem assumes identity fp16 rounding; inspect active-head and positive-size hypotheses. | Statement and preconditionsspecification attention_forward_triton_closed_form_correct
(Q K V Q_scale K_scale Out : RegionName) (s : BlockState)
(stride_qz stride_qh Z H BLOCK_M BLOCK_N numKVBlocks
HEAD_DIM BLOCK_DMODEL HEAD_ACTIVE STAGE : Nat)
(hBN : 0 < BLOCK_N) (hActiveLe : HEAD_ACTIVE ≤ BLOCK_DMODEL)
(hHD : HEAD_ACTIVE ≤ HEAD_DIM) (hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_forward_triton_surface Q K V Q_scale K_scale Out
stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1
Z H (BLOCK_N * numKVBlocks) HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL
HEAD_ACTIVE STAGE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
active s (BLOCK_N * numKVBlocks) HEAD_ACTIVE BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => (Out,
outOffset s H stride_qz stride_qh HEAD_DIM 1 BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
if h : idx.2.1.val < HEAD_ACTIVE then
attentionRealBase2PerKeyScale
(qTile s Q H stride_qz stride_qh HEAD_DIM BLOCK_M HEAD_ACTIVE)
(kTile s K H stride_qz stride_qh HEAD_DIM (BLOCK_N * numKVBlocks) HEAD_ACTIVE)
(vTile s V H stride_qz stride_qh HEAD_DIM (BLOCK_N * numKVBlocks) HEAD_ACTIVE)
(keyScale s Q_scale K_scale (BLOCK_N * numKVBlocks) BLOCK_M BLOCK_N
(BLOCK_N * numKVBlocks))
(idx.1, ⟨idx.2.1.val, h⟩, PUnit.unit)
else (0 : ℝ))
Referenced model/specification definitions: Python source |
|---|
attention_forward_tritonattention_forward_triton_io_correctnessAbstract cast/store rounding | Configured model / stage Forward attention with the explicitly shared contiguous head strides and a whole number of KV blocks. The IO theorem assumes identity fp16 rounding; inspect active-head and positive-size hypotheses. | Statement and preconditionsspecification attention_forward_triton_io_correctness (R : RoundingModel)
(hfp16 : R.round .fp16 = id)
(Q K V QScale KScale Out : RegionName)
(stride_qz stride_qh Z H HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE
numKVBlocks : Nat)
(hBD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N) (hnum : 0 < numKVBlocks)
(hActiveLe : HEAD_ACTIVE ≤ BLOCK_DMODEL) (hHD : HEAD_ACTIVE ≤ HEAD_DIM) :
attnFwdIO Q K V QScale KScale Out stride_qz stride_qh Z H HEAD_DIM BLOCK_M BLOCK_N
BLOCK_DMODEL HEAD_ACTIVE STAGE numKVBlocks ⊨[R]
fun p₀ _ _ xs ys zs ws vs j =>
if h : j.val % BLOCK_DMODEL < HEAD_ACTIVE then
attnFwdIOOutSpec p₀ BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE numKVBlocks xs ys zs ws vs
((Lane2D.decode j).1, ⟨j.val % BLOCK_DMODEL, h⟩, PUnit.unit)
else 0
IO definition: attnFwdIOdef attnFwdIO (Q K V QScale KScale Out : RegionName)
(stride_qz stride_qh Z H HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE
numKVBlocks : Nat) : StreamMasked3DKernelIO₅ where
kernel := attention_forward_triton_surface Q K V QScale KScale Out
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
Z H (BLOCK_N * numKVBlocks) HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE
inp1 := Q
inp2 := K
inp3 := V
inp4 := QScale
inp5 := KScale
out := Out
T := numKVBlocks
B1 := BLOCK_M * BLOCK_DMODEL
B2 := BLOCK_DMODEL * BLOCK_N
B3 := BLOCK_N * BLOCK_DMODEL
B4 := 1
B5 := 1
C := BLOCK_M * BLOCK_DMODEL
outDType := .real
read1 := fun p₀ p₁ _ _ j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (p₀ * BLOCK_M + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
read2 := fun _ p₁ _ t j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ j.val / BLOCK_N + (t.val * BLOCK_N + j.val % BLOCK_N) * HEAD_DIM
read3 := fun _ p₁ _ t j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (t.val * BLOCK_N + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
read4 := fun p₀ p₁ _ _ _ =>
p₁ * ((BLOCK_N * numKVBlocks + BLOCK_M - 1) / BLOCK_M) + p₀
read5 := fun _ p₁ _ t _ =>
p₁ * ((BLOCK_N * numKVBlocks + BLOCK_N - 1) / BLOCK_N) + t.val
write := fun p₀ p₁ _ j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (p₀ * BLOCK_M + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
mask1 := fun p₀ _ _ _ j =>
p₀ * BLOCK_M + j.val / BLOCK_DMODEL < BLOCK_N * numKVBlocks
∧ j.val % BLOCK_DMODEL < HEAD_ACTIVE
mask2 := fun _ _ _ t j =>
j.val % BLOCK_N < BLOCK_N * numKVBlocks - t.val * BLOCK_N ∧ j.val / BLOCK_N < HEAD_ACTIVE
mask3 := fun _ _ _ t j =>
j.val / BLOCK_DMODEL < BLOCK_N * numKVBlocks - t.val * BLOCK_N
∧ j.val % BLOCK_DMODEL < HEAD_ACTIVE
mask4 := fun _ _ _ _ _ => True
mask5 := fun _ _ _ _ _ => True
writeMask := fun p₀ _ _ j =>
p₀ * BLOCK_M + j.val / BLOCK_DMODEL < BLOCK_N * numKVBlocks
∧ j.val % BLOCK_DMODEL < HEAD_ACTIVE
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton1attention_fwd_triton1_output_summary_generalMathematical execution | Configured model / stage The output contract uses STORE=false and IFCOND=false with T=NT*BT and contiguous head lanes. Projection facts for other flags are not output proofs for those flags. | Statement and preconditionsspecification attention_fwd_triton1_output_summary_general
(Q K V H O : RegionName) (s_qh s_hh s_ht : Nat) (scale : ℝ)
(BT BD NT : Nat) (hBT : 0 < BT) (s : BlockState)
(hOQ : O ≠ Q) (hOK : O ≠ K) (hOV : O ≠ V) (hOH : O ≠ H) :
-- (1) all four STORE/IFCOND branch surfaces lower to the algorithm layer
((∃ alg, (attention_fwd_kernel_surface Q K V H O
s_qh BD 1 s_hh s_ht (NT * BT) scale BT BD NT Bool.false Bool.false).toAlgorithm?
= Except.ok alg) ∧
(∃ alg, (attention_fwd_kernel_surface Q K V H O
s_qh BD 1 s_hh s_ht (NT * BT) scale BT BD NT Bool.true Bool.false).toAlgorithm?
= Except.ok alg) ∧
(∃ alg, (attention_fwd_kernel_surface Q K V H O
s_qh BD 1 s_hh s_ht (NT * BT) scale BT BD NT Bool.false Bool.true).toAlgorithm?
= Except.ok alg) ∧
(∃ alg, (attention_fwd_kernel_surface Q K V H O
s_qh BD 1 s_hh s_ht (NT * BT) scale BT BD NT Bool.true Bool.true).toAlgorithm?
= Except.ok alg)) ∧
-- (2) the default branch runs to completion (existence / termination)
(∃ sF, exec (attention_fwd_kernel_surface Q K V H O
s_qh BD 1 s_hh s_ht (NT * BT) scale BT BD NT Bool.false Bool.false).toAlgKernel s
= some sF) ∧
-- (3) standard Realizes_without_Rounding: every streamed O lane holds the genuine closed form
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_fwd_kernel_surface Q K V H O
s_qh BD 1 s_hh s_ht (NT * BT) scale BT BD NT Bool.false Bool.false)
(initialState := s)
(write := fun i : Fin NT × Fin BT × Fin BD =>
some (O, s.pids 0 * s_qh + (i.1.val * BT + i.2.1.val) * BD + i.2.2.val))
(expected := fun i : Fin NT × Fin BT × Fin BD =>
outputClosedForm s Q K V scale BT BD
(aft1QAddrG s s_qh BT BD) (aft1KAddrG s s_qh BT BD)
(aft1QAddrG s s_qh BT BD) i.1.val i.2.1 i.2.2)
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton1attention_fwd_triton1_io_correctnessAbstract cast/store rounding | Configured model / stage The output contract uses STORE=false and IFCOND=false with T=NT*BT and contiguous head lanes. Projection facts for other flags are not output proofs for those flags. | Statement and preconditionsspecification attention_fwd_triton1_io_correctness (R : RoundingModel)
(Q K V H O : RegionName) (s_qh s_hh s_ht : Nat) (scale : ℝ)
(BT BD NT : Nat) (hBT : 0 < BT)
(hOQ : O ≠ Q) (hOK : O ≠ K) (hOV : O ≠ V) (hOH : O ≠ H) :
attentionFwdTriton1KernelIO Q K V H O s_qh s_hh s_ht scale BT BD NT ⊨[R]
fun _ _ qs ks vs t j => aft1StreamSpec scale NT BT BD qs ks vs t j
IO definition: attentionFwdTriton1KernelIOdef attentionFwdTriton1KernelIO (Q K V H O : RegionName)
(s_qh s_hh s_ht : Nat) (scale : ℝ) (BT BD NT : Nat) :
StreamEmitMasked2DKernelIO₃ where
kernel := attention_fwd_kernel_surface Q K V H O s_qh BD 1 s_hh s_ht (NT * BT)
scale BT BD NT Bool.false Bool.false
inp1 := Q
inp2 := K
inp3 := V
out := O
T := NT
B1 := BT * BD
B2 := BD * BT
B3 := BT * BD
C := BT * BD
read1 := fun p₀ _ t j => p₀ * s_qh + (t.val * BT + j.val / BD) * BD + j.val % BD
read2 := fun p₀ _ t j => p₀ * s_qh + j.val / BT + (t.val * BT + j.val % BT) * BD
read3 := fun p₀ _ t j => p₀ * s_qh + (t.val * BT + j.val / BD) * BD + j.val % BD
write := fun p₀ _ t j => p₀ * s_qh + (t.val * BT + j.val / BD) * BD + j.val % BD
mask1 := fun _ _ _ _ => True
mask2 := fun _ _ _ _ => True
mask3 := fun _ _ _ _ => True
writeMask := fun _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton2attention_fwd_triton2_output_summary_generalMathematical execution | Configured model / stage Base-2 scaled attention with shared contiguous head strides and N_CTX=BLOCK_N*numKVBlocks. The IO theorem requires identity fp16 rounding. | Statement and preconditionsspecification attention_fwd_triton2_output_summary_general
(Q K V Q_scale K_scale Out : RegionName) (s : BlockState)
(stride_qz stride_qh Z H BLOCK_M BLOCK_N numKVBlocks
HEAD_DIM BLOCK_DMODEL HEAD_ACTIVE STAGE : Nat)
(hBN : 0 < BLOCK_N) (hActiveLe : HEAD_ACTIVE ≤ BLOCK_DMODEL)
(hHD : HEAD_ACTIVE ≤ HEAD_DIM) (hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (attention_fwd_triton2_surface Q K V Q_scale K_scale Out
stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1
Z H (BLOCK_N * numKVBlocks) HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL
HEAD_ACTIVE STAGE).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_fwd_triton2_surface Q K V Q_scale K_scale Out
stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1
Z H (BLOCK_N * numKVBlocks) HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL
HEAD_ACTIVE STAGE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
active s (BLOCK_N * numKVBlocks) HEAD_ACTIVE BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => (Out,
outOffset s H stride_qz stride_qh HEAD_DIM 1 BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
if h : idx.2.1.val < HEAD_ACTIVE then
attentionRealBase2PerKeyScale
(qTile s Q H stride_qz stride_qh HEAD_DIM BLOCK_M HEAD_ACTIVE)
(kTile s K H stride_qz stride_qh HEAD_DIM (BLOCK_N * numKVBlocks) HEAD_ACTIVE)
(vTile s V H stride_qz stride_qh HEAD_DIM (BLOCK_N * numKVBlocks) HEAD_ACTIVE)
(keyScale s Q_scale K_scale (BLOCK_N * numKVBlocks) BLOCK_M BLOCK_N
(BLOCK_N * numKVBlocks))
(idx.1, ⟨idx.2.1.val, h⟩, PUnit.unit)
else (0 : ℝ))
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton2attention_fwd_triton2_io_correctnessAbstract cast/store rounding | Configured model / stage Base-2 scaled attention with shared contiguous head strides and N_CTX=BLOCK_N*numKVBlocks. The IO theorem requires identity fp16 rounding. | Statement and preconditionsspecification attention_fwd_triton2_io_correctness (R : RoundingModel)
(hfp16 : R.round .fp16 = id)
(Q K V Q_scale K_scale Out : RegionName)
(stride_qz stride_qh Z H HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE
numKVBlocks : Nat)
(hBN : 0 < BLOCK_N) (hnum : 0 < numKVBlocks)
(hActiveLe : HEAD_ACTIVE ≤ BLOCK_DMODEL) (hBDHD : BLOCK_DMODEL ≤ HEAD_DIM) :
attentionFwdTriton2IO Q K V Q_scale K_scale Out stride_qz stride_qh Z H HEAD_DIM
BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE numKVBlocks ⊨[R]
fun _ _ _ xs ys zs ws vs j =>
attentionFwdTriton2IOOutSpec BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE numKVBlocks
xs ys zs ws vs (Lane2D.decode j)
IO definition: attentionFwdTriton2IOdef attentionFwdTriton2IO (Q K V Q_scale K_scale Out : RegionName)
(stride_qz stride_qh Z H HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE
numKVBlocks : Nat) : StreamMasked3DKernelIO₅ where
kernel := attention_fwd_triton2_surface Q K V Q_scale K_scale Out
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
Z H (BLOCK_N * numKVBlocks) HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE
inp1 := Q
inp2 := K
inp3 := V
inp4 := Q_scale
inp5 := K_scale
out := Out
T := numKVBlocks
B1 := BLOCK_M * BLOCK_DMODEL
B2 := BLOCK_DMODEL * BLOCK_N
B3 := BLOCK_N * BLOCK_DMODEL
B4 := 1
B5 := 1
C := BLOCK_M * BLOCK_DMODEL
outDType := .real
read1 := fun p₀ p₁ _ _ j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (p₀ * BLOCK_M + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
read2 := fun _ p₁ _ t j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ j.val / BLOCK_N + (t.val * BLOCK_N + j.val % BLOCK_N) * HEAD_DIM
read3 := fun _ p₁ _ t j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (t.val * BLOCK_N + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
read4 := fun p₀ p₁ _ _ _ => p₁ * cdiv (BLOCK_N * numKVBlocks) BLOCK_M + p₀
read5 := fun _ p₁ _ t _ => p₁ * cdiv (BLOCK_N * numKVBlocks) BLOCK_N + t.val
write := fun p₀ p₁ _ j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (p₀ * BLOCK_M + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
mask1 := fun p₀ _ _ _ j =>
p₀ * BLOCK_M + j.val / BLOCK_DMODEL < BLOCK_N * numKVBlocks ∧
j.val % BLOCK_DMODEL < HEAD_ACTIVE
mask2 := fun _ _ _ t j =>
j.val % BLOCK_N < BLOCK_N * numKVBlocks - t.val * BLOCK_N ∧ j.val / BLOCK_N < HEAD_ACTIVE
mask3 := fun _ _ _ t j =>
j.val / BLOCK_DMODEL < BLOCK_N * numKVBlocks - t.val * BLOCK_N ∧
j.val % BLOCK_DMODEL < HEAD_ACTIVE
mask4 := fun _ _ _ _ _ => True
mask5 := fun _ _ _ _ _ => True
writeMask := fun p₀ _ _ j =>
p₀ * BLOCK_M + j.val / BLOCK_DMODEL < BLOCK_N * numKVBlocks ∧
j.val % BLOCK_DMODEL < HEAD_ACTIVE
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton3attention_fwd_triton3_python_case1_output_summary_generalMathematical execution | Configured model / stage Four explicitly configured attention branches, including sliding-window, complementary-window, and continuation cases. Case-specific O/L/M targets, divisibility, and rounding assumptions are in each statement. | Statement and preconditionsspecification attention_fwd_triton3_python_case1_output_summary_general
(Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN : Nat) (s : BlockState)
(hND : 0 < ND) (hBM : 0 < BM) (hBN : 0 < BN) (hNC : 0 < NKV_CTX) (hBNdvd : BN ∣ NKV_CTX)
(hH : 0 < H) (hHKV : H_KV = H)
(hskz : skz = sqz) (hskh : skh = sqh) (hsvz : svz = sqz) (hsvh : svh = sqh)
(hsoz : soz = sqz) (hsoh : soh = sqh)
(hMO : M ≠ Out) (hundef : ∀ rg o, s.undef rg o = 0)
(hinjO : Function.Injective
(fun idx : TileIndex [BM, ND] => (s.pids 1 / H * sqz + s.pids 1 % H * sqh) + (s.pids 0 * BM + idx.1.val) * som + idx.2.1.val * son))
(hinjM : Function.Injective
(fun r : TileIndex [BM] => s.pids 1 * ROUND_CTX + (s.pids 0 * BM + r.1.val))) :
(∃ alg, (attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 1 1 0).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 1 1 0)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BM, ND] => active s N_CTX ND BM idx)
(fun idx : TileIndex [BM, ND] => (Out, outOffset s H sqz sqh som son BM idx)))
(expected := fun idx : TileIndex [BM, ND] =>
attentionFwdTriton3Case1OutSpecG s Q K V (s.pids 1 / H * sqz + s.pids 1 % H * sqh) BM ND NKV_CTX sqm sqk skn skk svk svn (sm_scale * 1.4426950408889634) BN off size idx) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 1 1 0)
(initialState := s)
(write := fun i : Fin BM => some (M, lRowOffset s (s.pids 1) ROUND_CTX BM i))
(expected := fun i : Fin BM =>
attentionFwdTriton3KMSpecG s Q K V (s.pids 1 / H * sqz + s.pids 1 % H * sqh) BM ND NKV_CTX sqm sqk skn skk svk svn (sm_scale * 1.4426950408889634) (fun i j => natSlidingWindowKeepG (s.pids 0) BM BN off size i j) i hND)
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton3attention_fwd_triton3_python_case2_output_summary_generalMathematical execution | Configured model / stage Four explicitly configured attention branches, including sliding-window, complementary-window, and continuation cases. Case-specific O/L/M targets, divisibility, and rounding assumptions are in each statement. | Statement and preconditionsspecification attention_fwd_triton3_python_case2_output_summary_general
(Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN : Nat) (s : BlockState)
(hND : 0 < ND) (hBM : 0 < BM) (hBN : 0 < BN) (hNC : 0 < NKV_CTX) (hBNdvd : BN ∣ NKV_CTX)
(hH : 0 < H) (hHKV : H_KV = H)
(hskz : skz = sqz) (hskh : skh = sqh) (hsvz : svz = sqz) (hsvh : svh = sqh)
(hsoz : soz = sqz) (hsoh : soh = sqh)
(hMO : M ≠ Out) (hundef : ∀ rg o, s.undef rg o = 0)
(hinjO : Function.Injective
(fun idx : TileIndex [BM, ND] => (s.pids 1 / H * sqz + s.pids 1 % H * sqh) + (s.pids 0 * BM + idx.1.val) * som + idx.2.1.val * son))
(hinjM : Function.Injective
(fun r : TileIndex [BM] => s.pids 1 * ROUND_CTX + (s.pids 0 * BM + r.1.val))) :
(∃ alg, (attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 1 1 1).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 1 1 1)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BM, ND] => active s N_CTX ND BM idx)
(fun idx : TileIndex [BM, ND] => (Out, outOffset s H sqz sqh som son BM idx)))
(expected := fun idx : TileIndex [BM, ND] =>
attentionFwdTriton3Case2OutSpecG s Q K V (s.pids 1 / H * sqz + s.pids 1 % H * sqh) BM ND NKV_CTX sqm sqk skn skk svk svn (sm_scale * 1.4426950408889634) BN off size idx) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 1 1 1)
(initialState := s)
(write := fun i : Fin BM => some (M, lRowOffset s (s.pids 1) ROUND_CTX BM i))
(expected := fun i : Fin BM =>
attentionFwdTriton3KMSpecG s Q K V (s.pids 1 / H * sqz + s.pids 1 % H * sqh) BM ND NKV_CTX sqm sqk skn skk svk svn (sm_scale * 1.4426950408889634) (fun i j => natComplementSlidingWindowKeepG (s.pids 0) BM BN off size i j) i hND)
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton3attention_fwd_triton3_python_case3_output_summary_generalMathematical execution | Configured model / stage Four explicitly configured attention branches, including sliding-window, complementary-window, and continuation cases. Case-specific O/L/M targets, divisibility, and rounding assumptions are in each statement. | Statement and preconditionsspecification attention_fwd_triton3_python_case3_output_summary_general
(Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off BM ND BN : Nat) (s : BlockState)
(hND : 0 < ND) (hBM : 0 < BM) (hBN : 0 < BN) (hNC : 0 < NKV_CTX) (hBNdvd : BN ∣ NKV_CTX)
(hH : 0 < H) (hHKV : H_KV = H)
(hskz : skz = sqz) (hskh : skh = sqh) (hsvz : svz = sqz) (hsvh : svh = sqh)
(hsoz : soz = sqz) (hsoh : soh = sqh)
(hMO : M ≠ Out) (hundef : ∀ rg o, s.undef rg o = 0)
(hinjO : Function.Injective
(fun idx : TileIndex [BM, ND] => (s.pids 1 / H * sqz + s.pids 1 % H * sqh) + (s.pids 0 * BM + idx.1.val) * som + idx.2.1.val * son))
(hinjM : Function.Injective
(fun r : TileIndex [BM] => s.pids 1 * ROUND_CTX + (s.pids 0 * BM + r.1.val))) :
(∃ alg, (attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off 0 1 1 BM ND BN 1 1 0 0).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off 0 1 1 BM ND BN 1 1 0 0)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BM, ND] => active s N_CTX ND BM idx)
(fun idx : TileIndex [BM, ND] => (Out, outOffset s H sqz sqh som son BM idx)))
(expected := fun idx : TileIndex [BM, ND] =>
attentionFwdTriton3Case3OutSpecG s Q K V (s.pids 1 / H * sqz + s.pids 1 % H * sqh) BM ND NKV_CTX sqm sqk skn skk svk svn (sm_scale * 1.4426950408889634) idx) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off 0 1 1 BM ND BN 1 1 0 0)
(initialState := s)
(write := fun i : Fin BM => some (M, lRowOffset s (s.pids 1) ROUND_CTX BM i))
(expected := fun i : Fin BM =>
attentionFwdTriton3Case3MSpecG s Q K V (s.pids 1 / H * sqz + s.pids 1 % H * sqh) BM ND NKV_CTX sqm sqk skn skk svk svn (sm_scale * 1.4426950408889634) i hND)
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton3attention_fwd_triton3_python_case4_output_summary_generalMathematical execution | Configured model / stage Four explicitly configured attention branches, including sliding-window, complementary-window, and continuation cases. Case-specific O/L/M targets, divisibility, and rounding assumptions are in each statement. | Statement and preconditionsspecification attention_fwd_triton3_python_case4_output_summary_general
(Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN : Nat) (s : BlockState)
(hND : 0 < ND) (hBM : 0 < BM) (hBN : 0 < BN) (hNC : 0 < NKV_CTX) (hBNdvd : BN ∣ NKV_CTX)
(hH : 0 < H) (hHKV : H_KV = H)
(hskz : skz = sqz) (hskh : skh = sqh) (hsvz : svz = sqz) (hsvh : svh = sqh)
(hsoz : soz = sqz) (hsoh : soh = sqh)
(hMO : M ≠ Out) (hundef : ∀ rg o, s.undef rg o = 0)
(hinjO : Function.Injective
(fun idx : TileIndex [BM, ND] => (s.pids 1 / H * sqz + s.pids 1 % H * sqh) + (s.pids 0 * BM + idx.1.val) * som + idx.2.1.val * son))
(hinjM : Function.Injective
(fun r : TileIndex [BM] => s.pids 1 * ROUND_CTX + (s.pids 0 * BM + r.1.val))) :
(∃ alg, (attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 0 1 0).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 0 1 0)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BM, ND] => active s N_CTX ND BM idx)
(fun idx : TileIndex [BM, ND] => (Out, outOffset s H sqz sqh som son BM idx)))
(expected := fun idx : TileIndex [BM, ND] =>
attentionFwdTriton3Case4OutSpecG s Q K V M Out L (s.pids 1 / H * sqz + s.pids 1 % H * sqh) BM ND NKV_CTX sqm sqk skn skk svk svn som son ROUND_CTX (sm_scale * 1.4426950408889634) BN off size hND idx) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 0 1 0)
(initialState := s)
(write := fun i : Fin BM => some (M, lRowOffset s (s.pids 1) ROUND_CTX BM i))
(expected := fun i : Fin BM =>
attentionFwdTriton3Case4MSpecG s Q K V M Out L (s.pids 1 / H * sqz + s.pids 1 % H * sqh) BM ND NKV_CTX sqm sqk skn skk svk svn som son ROUND_CTX (sm_scale * 1.4426950408889634) BN off size i hND)
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton3attention_fwd_triton3_case3_io_correctnessAbstract cast/store rounding | Configured model / stage Four explicitly configured attention branches, including sliding-window, complementary-window, and continuation cases. Case-specific O/L/M targets, divisibility, and rounding assumptions are in each statement. | Statement and preconditionsspecification attention_fwd_triton3_case3_io_correctness (R : RoundingModel)
(Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off BM ND BN : Nat)
(hND : 0 < ND) (hBM : 0 < BM) (hBN : 0 < BN) (hNC : 0 < NKV_CTX) (hBNdvd : BN ∣ NKV_CTX)
(hH : 0 < H) (hHKV : H_KV = H)
(hskz : skz = sqz) (hskh : skh = sqh) (hsvz : svz = sqz) (hsvh : svh = sqh)
(hsoz : soz = sqz) (hsoh : soh = sqh)
(hMO : M ≠ Out)
(hinjO : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BM, ND] =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + idx.1.val) * som + idx.2.1.val * son))
(hinjM : ∀ p₀ p₁ : Nat, Function.Injective
(fun r : TileIndex [BM] => p₁ * ROUND_CTX + (p₀ * BM + r.1.val))) :
attentionFwdTriton3Case3IO Q K V M Out L sm_scale sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son Z H H_KV N_CTX ROUND_CTX NKV_CTX off BM ND BN ⊨[R]
fun _ _ _ xs ys zs =>
(fun j => attentionFwdTriton3Case3IOOutSpec BM ND BN NKV_CTX (NKV_CTX / BN)
(Nat.div_pos (Nat.le_of_dvd hNC hBNdvd) hBN) (Nat.div_mul_cancel hBNdvd) hBN
(sm_scale * 1.4426950408889634) xs ys zs j,
fun i => attentionFwdTriton3Case3IOMSpec BM ND BN NKV_CTX (NKV_CTX / BN) hND
(Nat.div_pos (Nat.le_of_dvd hNC hBNdvd) hBN) (Nat.div_mul_cancel hBNdvd) hBN
(sm_scale * 1.4426950408889634) xs ys zs i)
IO definition: attentionFwdTriton3Case3IOdef attentionFwdTriton3Case3IO (Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off BM ND BN : Nat) :
StreamMasked3DKernelIO₃ₓ₂ where
kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off 0 1 1 BM ND BN 1 1 0 0
inp1 := Q
inp2 := K
inp3 := V
out1 := Out
out2 := M
T := NKV_CTX / BN
B1 := BM * ND
B2 := ND * BN
B3 := BN * ND
C1 := BM * ND
C2 := BM
read1 := fun p₀ p₁ _ _ j =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + j.val / ND) * sqm + (j.val % ND) * sqk
read2 := fun _ p₁ _ t j =>
(p₁ / H * sqz + p₁ % H * sqh) + (j.val / BN) * skk + (t.val * BN + j.val % BN) * skn
read3 := fun _ p₁ _ t j =>
(p₁ / H * sqz + p₁ % H * sqh) + (t.val * BN + j.val / ND) * svk + (j.val % ND) * svn
write1 := fun p₀ p₁ _ j =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + j.val / ND) * som + (j.val % ND) * son
write2 := fun p₀ p₁ _ j => p₁ * ROUND_CTX + (p₀ * BM + j.val)
mask1 := fun _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ => True
mask3 := fun _ _ _ _ _ => True
writeMask1 := fun _ _ _ _ => True
writeMask2 := fun _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton3attention_fwd_triton3_case1_io_correctnessAbstract cast/store rounding | Configured model / stage Four explicitly configured attention branches, including sliding-window, complementary-window, and continuation cases. Case-specific O/L/M targets, divisibility, and rounding assumptions are in each statement. | Statement and preconditionsspecification attention_fwd_triton3_case1_io_correctness (R : RoundingModel)
(Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN : Nat)
(hND : 0 < ND) (hBM : 0 < BM) (hBN : 0 < BN) (hNC : 0 < NKV_CTX) (hBNdvd : BN ∣ NKV_CTX)
(hH : 0 < H) (hHKV : H_KV = H)
(hskz : skz = sqz) (hskh : skh = sqh) (hsvz : svz = sqz) (hsvh : svh = sqh)
(hsoz : soz = sqz) (hsoh : soh = sqh)
(hMO : M ≠ Out)
(hinjO : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BM, ND] =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + idx.1.val) * som + idx.2.1.val * son))
(hinjM : ∀ p₀ p₁ : Nat, Function.Injective
(fun r : TileIndex [BM] => p₁ * ROUND_CTX + (p₀ * BM + r.1.val))) :
attentionFwdTriton3Case1IO Q K V M Out L sm_scale sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN ⊨[R]
fun p₀ _ _ xs ys zs =>
(fun j => attentionFwdTriton3Case1IOOutSpec p₀ BM ND BN NKV_CTX (NKV_CTX / BN)
(Nat.div_pos (Nat.le_of_dvd hNC hBNdvd) hBN) (Nat.div_mul_cancel hBNdvd) hBN
(sm_scale * 1.4426950408889634) off size xs ys zs j,
fun i => attentionFwdTriton3Case1IOMSpec p₀ BM ND BN NKV_CTX (NKV_CTX / BN) hND
(Nat.div_pos (Nat.le_of_dvd hNC hBNdvd) hBN) (Nat.div_mul_cancel hBNdvd) hBN
(sm_scale * 1.4426950408889634) off size xs ys zs i)
IO definition: attentionFwdTriton3Case1IOdef attentionFwdTriton3Case1IO (Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN : Nat) :
StreamMasked3DKernelIO₃ₓ₂ where
kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 1 1 0
inp1 := Q
inp2 := K
inp3 := V
out1 := Out
out2 := M
T := NKV_CTX / BN
B1 := BM * ND
B2 := ND * BN
B3 := BN * ND
C1 := BM * ND
C2 := BM
read1 := fun p₀ p₁ _ _ j =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + j.val / ND) * sqm + (j.val % ND) * sqk
read2 := fun _ p₁ _ t j =>
(p₁ / H * sqz + p₁ % H * sqh) + (j.val / BN) * skk + (t.val * BN + j.val % BN) * skn
read3 := fun _ p₁ _ t j =>
(p₁ / H * sqz + p₁ % H * sqh) + (t.val * BN + j.val / ND) * svk + (j.val % ND) * svn
write1 := fun p₀ p₁ _ j =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + j.val / ND) * som + (j.val % ND) * son
write2 := fun p₀ p₁ _ j => p₁ * ROUND_CTX + (p₀ * BM + j.val)
mask1 := fun _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ => True
mask3 := fun _ _ _ _ _ => True
writeMask1 := fun _ _ _ _ => True
writeMask2 := fun _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton3attention_fwd_triton3_case2_io_correctnessAbstract cast/store rounding | Configured model / stage Four explicitly configured attention branches, including sliding-window, complementary-window, and continuation cases. Case-specific O/L/M targets, divisibility, and rounding assumptions are in each statement. | Statement and preconditionsspecification attention_fwd_triton3_case2_io_correctness (R : RoundingModel)
(Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN : Nat)
(hND : 0 < ND) (hBM : 0 < BM) (hBN : 0 < BN) (hNC : 0 < NKV_CTX) (hBNdvd : BN ∣ NKV_CTX)
(hH : 0 < H) (hHKV : H_KV = H)
(hskz : skz = sqz) (hskh : skh = sqh) (hsvz : svz = sqz) (hsvh : svh = sqh)
(hsoz : soz = sqz) (hsoh : soh = sqh)
(hMO : M ≠ Out)
(hinjO : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BM, ND] =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + idx.1.val) * som + idx.2.1.val * son))
(hinjM : ∀ p₀ p₁ : Nat, Function.Injective
(fun r : TileIndex [BM] => p₁ * ROUND_CTX + (p₀ * BM + r.1.val))) :
attentionFwdTriton3Case2IO Q K V M Out L sm_scale sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN ⊨[R]
fun p₀ _ _ xs ys zs =>
(fun j => attentionFwdTriton3Case2IOOutSpec p₀ BM ND BN NKV_CTX (NKV_CTX / BN)
(Nat.div_pos (Nat.le_of_dvd hNC hBNdvd) hBN) (Nat.div_mul_cancel hBNdvd) hBN
(sm_scale * 1.4426950408889634) off size xs ys zs j,
fun i => attentionFwdTriton3Case2IOMSpec p₀ BM ND BN NKV_CTX (NKV_CTX / BN) hND
(Nat.div_pos (Nat.le_of_dvd hNC hBNdvd) hBN) (Nat.div_mul_cancel hBNdvd) hBN
(sm_scale * 1.4426950408889634) off size xs ys zs i)
IO definition: attentionFwdTriton3Case2IOdef attentionFwdTriton3Case2IO (Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN : Nat) :
StreamMasked3DKernelIO₃ₓ₂ where
kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 1 1 1
inp1 := Q
inp2 := K
inp3 := V
out1 := Out
out2 := M
T := NKV_CTX / BN
B1 := BM * ND
B2 := ND * BN
B3 := BN * ND
C1 := BM * ND
C2 := BM
read1 := fun p₀ p₁ _ _ j =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + j.val / ND) * sqm + (j.val % ND) * sqk
read2 := fun _ p₁ _ t j =>
(p₁ / H * sqz + p₁ % H * sqh) + (j.val / BN) * skk + (t.val * BN + j.val % BN) * skn
read3 := fun _ p₁ _ t j =>
(p₁ / H * sqz + p₁ % H * sqh) + (t.val * BN + j.val / ND) * svk + (j.val % ND) * svn
write1 := fun p₀ p₁ _ j =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + j.val / ND) * som + (j.val % ND) * son
write2 := fun p₀ p₁ _ j => p₁ * ROUND_CTX + (p₀ * BM + j.val)
mask1 := fun _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ => True
mask3 := fun _ _ _ _ _ => True
writeMask1 := fun _ _ _ _ => True
writeMask2 := fun _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
attention_fwd_triton3attention_fwd_triton3_case4_io_correctnessAbstract cast/store rounding | Configured model / stage Four explicitly configured attention branches, including sliding-window, complementary-window, and continuation cases. Case-specific O/L/M targets, divisibility, and rounding assumptions are in each statement. | Statement and preconditionsspecification attention_fwd_triton3_case4_io_correctness (R : RoundingModel)
(Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN : Nat)
(hND : 0 < ND) (hBM : 0 < BM) (hBN : 0 < BN) (hNC : 0 < NKV_CTX) (hBNdvd : BN ∣ NKV_CTX)
(hH : 0 < H) (hHKV : H_KV = H)
(hskz : skz = sqz) (hskh : skh = sqh) (hsvz : svz = sqz) (hsvh : svh = sqh)
(hsoz : soz = sqz) (hsoh : soh = sqh)
(hMO : M ≠ Out)
(hinjO : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BM, ND] =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + idx.1.val) * som + idx.2.1.val * son))
(hinjM : ∀ p₀ p₁ : Nat, Function.Injective
(fun r : TileIndex [BM] => p₁ * ROUND_CTX + (p₀ * BM + r.1.val))) :
attentionFwdTriton3Case4IO Q K V M Out L sm_scale sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN ⊨[R]
fun p₀ _ _ x1s x2s x3s x4s x5s x6s =>
(fun j => attentionFwdTriton3Case4IOOutSpec p₀ BM ND BN NKV_CTX (NKV_CTX / BN) hND
(Nat.div_pos (Nat.le_of_dvd hNC hBNdvd) hBN) (Nat.div_mul_cancel hBNdvd) hBN
(sm_scale * 1.4426950408889634) off size x1s x2s x3s x4s x5s x6s j,
fun i => attentionFwdTriton3Case4IOMSpec p₀ BM ND BN NKV_CTX (NKV_CTX / BN) hND
(Nat.div_pos (Nat.le_of_dvd hNC hBNdvd) hBN) (Nat.div_mul_cancel hBNdvd) hBN
(sm_scale * 1.4426950408889634) off size x1s x2s x3s x4s x5s x6s i)
IO definition: attentionFwdTriton3Case4IOdef attentionFwdTriton3Case4IO (Q K V M Out L : RegionName) (sm_scale : ℝ)
(sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size BM ND BN : Nat) :
StreamMasked3DKernelIO₆ₓ₂ where
kernel := attention_fwd_triton3_surface Q K V M Out L sm_scale
sqz sqh sqm sqk skz skh skn skk svz svh svk svn soz soh som son
Z H H_KV N_CTX ROUND_CTX NKV_CTX off size 1 1 BM ND BN 1 0 1 0
inp1 := Q
inp2 := K
inp3 := V
inp4 := M
inp5 := L
inp6 := Out
out1 := Out
out2 := M
T := NKV_CTX / BN
B1 := BM * ND
B2 := ND * BN
B3 := BN * ND
B4 := BM
B5 := BM
B6 := BM * ND
C1 := BM * ND
C2 := BM
read1 := fun p₀ p₁ _ _ j =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + j.val / ND) * sqm + (j.val % ND) * sqk
read2 := fun _ p₁ _ t j =>
(p₁ / H * sqz + p₁ % H * sqh) + (j.val / BN) * skk + (t.val * BN + j.val % BN) * skn
read3 := fun _ p₁ _ t j =>
(p₁ / H * sqz + p₁ % H * sqh) + (t.val * BN + j.val / ND) * svk + (j.val % ND) * svn
read4 := fun p₀ p₁ _ _ j => p₁ * ROUND_CTX + (p₀ * BM + j.val)
read5 := fun p₀ p₁ _ _ j => p₁ * ROUND_CTX + (p₀ * BM + j.val)
read6 := fun p₀ p₁ _ _ j =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + j.val / ND) * som + (j.val % ND) * son
write1 := fun p₀ p₁ _ j =>
(p₁ / H * sqz + p₁ % H * sqh) + (p₀ * BM + j.val / ND) * som + (j.val % ND) * son
write2 := fun p₀ p₁ _ j => p₁ * ROUND_CTX + (p₀ * BM + j.val)
mask1 := fun _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ => True
mask3 := fun _ _ _ _ _ => True
mask4 := fun _ _ _ _ _ => True
mask5 := fun _ _ _ _ _ => True
mask6 := fun _ _ _ _ _ => True
writeMask1 := fun _ _ _ _ => True
writeMask2 := fun _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
attention_kernelattention_kernel_genuine_output_compute_correct_generalMathematical execution | Configured model / stage Aligned bias attention at the configured batch/head, bias, stride, and fp16 choices. The IO contract assumes identity fp16 rounding; no arbitrary-layout or dtype claim. | Statement and preconditionsspecification attention_kernel_genuine_output_compute_correct_general
(Q K V B0 Out : RegionName) (s : BlockState) (sm_scale : ℝ)
(stride_qh stride_kh stride_b0h BLOCK_M BLOCK_N HEAD BIAS_LAST_SIZE stride_b0m nB : Nat)
(hKN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M) (hHD : 0 < HEAD) (hnB : 1 ≤ nB)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_kernel_fwd_kernel_aligned_surface Q K V B0 Out sm_scale
stride_qh HEAD 1 stride_kh HEAD 1 stride_kh HEAD 1 stride_qh HEAD 1
stride_b0h stride_b0m 2 4 (BLOCK_N * nB) 0 BIAS_LAST_SIZE 128 HEAD BLOCK_M BLOCK_N
FloatDType.fp16)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, HEAD] =>
some (Out, surfaceOutOffset s stride_qh HEAD 1 BLOCK_M idx))
(expected := fun idx : TileIndex [BLOCK_M, HEAD] =>
MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (attentionKernelSpec s Q K V B0 sm_scale
(s.pids 1 * stride_qh) (s.pids 1 * stride_kh) (s.pids 1 * stride_b0h)
BLOCK_M BLOCK_N HEAD (BLOCK_N * nB) BIAS_LAST_SIZE stride_b0m (s.pids 0) idx))))
Referenced model/specification definitions: Python source |
|---|
attention_kernelattention_kernel_io_correctnessAbstract cast/store rounding | Configured model / stage Aligned bias attention at the configured batch/head, bias, stride, and fp16 choices. The IO contract assumes identity fp16 rounding; no arbitrary-layout or dtype claim. | Statement and preconditionsspecification attention_kernel_io_correctness (R : RoundingModel)
(hfp16 : R.round .fp16 = id)
(Q K V B0 Out : RegionName) (sm_scale : ℝ)
(stride_qh stride_kh stride_b0h BLOCK_M BLOCK_N HEAD BIAS_LAST_SIZE stride_b0m nB : Nat)
(hBM : 0 < BLOCK_M) (hBN : 0 < BLOCK_N) (hHD : 0 < HEAD) (hnB : 1 ≤ nB) :
attentionKernelIO Q K V B0 Out sm_scale stride_qh stride_kh stride_b0h
BLOCK_M BLOCK_N HEAD BIAS_LAST_SIZE stride_b0m nB ⊨[R]
fun _ _ _ xs ys zs ws j =>
attnGenScore (akIOscoreT nB BLOCK_M BLOCK_N HEAD sm_scale xs ys ws)
(akIOvT nB BLOCK_N HEAD zs) (Lane2D.decode j)
IO definition: attentionKernelIOdef attentionKernelIO (Q K V B0 Out : RegionName) (sm_scale : ℝ)
(stride_qh stride_kh stride_b0h BLOCK_M BLOCK_N HEAD BIAS_LAST_SIZE stride_b0m nB : Nat) :
StreamMasked3DKernelIO₄ where
kernel := attention_kernel_fwd_kernel_aligned_surface Q K V B0 Out sm_scale
stride_qh HEAD 1 stride_kh HEAD 1 stride_kh HEAD 1 stride_qh HEAD 1
stride_b0h stride_b0m 2 4 (BLOCK_N * nB) 0 BIAS_LAST_SIZE 128 HEAD BLOCK_M BLOCK_N
FloatDType.fp16
inp1 := Q
inp2 := K
inp3 := V
inp4 := B0
out := Out
T := nB
B1 := BLOCK_M * HEAD
B2 := HEAD * BLOCK_N
B3 := BLOCK_N * HEAD
B4 := BLOCK_M * (BLOCK_N + 1)
C := BLOCK_M * HEAD
outDType := .fp16
read1 := fun p₀ p₁ _ _ j =>
p₁ * stride_qh + (p₀ * BLOCK_M + j.val / HEAD) * HEAD + j.val % HEAD
read2 := fun _ p₁ _ t j =>
p₁ * stride_kh + j.val / BLOCK_N + (t.val * BLOCK_N + j.val % BLOCK_N) * HEAD
read3 := fun _ p₁ _ t j =>
p₁ * stride_kh + (t.val * BLOCK_N + j.val / HEAD) * HEAD + j.val % HEAD
read4 := fun p₀ p₁ _ t j =>
p₁ * stride_b0h + (p₀ * BLOCK_M + j.val / (BLOCK_N + 1)) * stride_b0m +
(if j.val % (BLOCK_N + 1) < BLOCK_N then
j.val % (BLOCK_N + 1) % BIAS_LAST_SIZE + BIAS_LAST_SIZE
else t.val)
write := fun p₀ p₁ _ j =>
p₁ * stride_qh + (p₀ * BLOCK_M + j.val / HEAD) * HEAD + (j.val % HEAD) * 1
mask1 := fun _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ => True
mask3 := fun _ _ _ _ _ => True
mask4 := fun _ _ _ _ _ => True
writeMask := fun _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
attention_kernel_alignedattention_kernel_aligned_output_summary_generalMathematical execution | Configured model / stage Aligned bias attention with shared contiguous Q/K/V/output layouts and fixed launch metadata. Identity fp16 rounding is required by the IO theorem. | Statement and preconditionsspecification attention_kernel_aligned_output_summary_general
(Q K V B0 Out : RegionName) (s : BlockState) (sm_scale : ℝ)
(stride_qh stride_b0h BLOCK_M BLOCK_N HEAD BIAS_LAST_SIZE stride_b0m nB : Nat)
(hKN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M) (hHD : 0 < HEAD) (hnB : 1 ≤ nB)
(hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (attention_kernel_aligned_fwd_kernel_aligned_surface Q K V B0 Out sm_scale
stride_qh HEAD 1 stride_qh HEAD 1 stride_qh HEAD 1 stride_qh HEAD 1
stride_b0h stride_b0m 2 4 (BLOCK_N * nB) 0 BIAS_LAST_SIZE 128 HEAD BLOCK_M BLOCK_N
FloatDType.fp16).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_kernel_aligned_fwd_kernel_aligned_surface Q K V B0 Out sm_scale
stride_qh HEAD 1 stride_qh HEAD 1 stride_qh HEAD 1 stride_qh HEAD 1
stride_b0h stride_b0m 2 4 (BLOCK_N * nB) 0 BIAS_LAST_SIZE 128 HEAD BLOCK_M BLOCK_N
FloatDType.fp16)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, HEAD] =>
some (Out, surfaceOutOffset s stride_qh HEAD 1 BLOCK_M idx))
(expected := fun idx : TileIndex [BLOCK_M, HEAD] =>
MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (alignedClosedForm s Q K V B0 sm_scale stride_qh stride_b0h stride_b0m
(BLOCK_N * nB) BIAS_LAST_SIZE HEAD BLOCK_M BLOCK_N idx))))
Referenced model/specification definitions: Python source |
|---|
attention_kernel_alignedattention_kernel_aligned_io_correctnessAbstract cast/store rounding | Configured model / stage Aligned bias attention with shared contiguous Q/K/V/output layouts and fixed launch metadata. Identity fp16 rounding is required by the IO theorem. | Statement and preconditionsspecification attention_kernel_aligned_io_correctness (R : RoundingModel)
(hfp16 : R.round .fp16 = id)
(Q K V B0 Out : RegionName) (sm_scale : ℝ)
(stride_qh stride_b0h BLOCK_M BLOCK_N HEAD BIAS_LAST_SIZE stride_b0m nB : Nat)
(hBM : 0 < BLOCK_M) (hBN : 0 < BLOCK_N) (hHD : 0 < HEAD) (hnB : 1 ≤ nB) :
attentionKernelAlignedIO Q K V B0 Out sm_scale stride_qh stride_b0h
BLOCK_M BLOCK_N HEAD BIAS_LAST_SIZE stride_b0m nB ⊨[R]
fun _ _ _ xs ys zs ws j =>
attnGenScore (akaIOscoreT nB BLOCK_M BLOCK_N HEAD sm_scale xs ys ws)
(akaIOvT nB BLOCK_N HEAD zs) (Lane2D.decode j)
IO definition: attentionKernelAlignedIOdef attentionKernelAlignedIO (Q K V B0 Out : RegionName) (sm_scale : ℝ)
(stride_qh stride_b0h BLOCK_M BLOCK_N HEAD BIAS_LAST_SIZE stride_b0m nB : Nat) :
StreamMasked3DKernelIO₄ where
kernel := attention_kernel_aligned_fwd_kernel_aligned_surface Q K V B0 Out sm_scale
stride_qh HEAD 1 stride_qh HEAD 1 stride_qh HEAD 1 stride_qh HEAD 1
stride_b0h stride_b0m 2 4 (BLOCK_N * nB) 0 BIAS_LAST_SIZE 128 HEAD BLOCK_M BLOCK_N
FloatDType.fp16
inp1 := Q
inp2 := K
inp3 := V
inp4 := B0
out := Out
T := nB
B1 := BLOCK_M * HEAD
B2 := HEAD * BLOCK_N
B3 := BLOCK_N * HEAD
B4 := BLOCK_M * (BLOCK_N + 1)
C := BLOCK_M * HEAD
outDType := .fp16
read1 := fun p₀ p₁ _ _ j =>
p₁ * stride_qh + (p₀ * BLOCK_M + j.val / HEAD) * HEAD + j.val % HEAD
read2 := fun _ p₁ _ t j =>
p₁ * stride_qh + j.val / BLOCK_N + (t.val * BLOCK_N + j.val % BLOCK_N) * HEAD
read3 := fun _ p₁ _ t j =>
p₁ * stride_qh + (t.val * BLOCK_N + j.val / HEAD) * HEAD + j.val % HEAD
read4 := fun p₀ p₁ _ t j =>
p₁ * stride_b0h + (p₀ * BLOCK_M + j.val / (BLOCK_N + 1)) * stride_b0m +
(if j.val % (BLOCK_N + 1) < BLOCK_N then
j.val % (BLOCK_N + 1) % BIAS_LAST_SIZE + BIAS_LAST_SIZE
else t.val)
write := fun p₀ p₁ _ j =>
p₁ * stride_qh + (p₀ * BLOCK_M + j.val / HEAD) * HEAD + (j.val % HEAD) * 1
mask1 := fun _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ => True
mask3 := fun _ _ _ _ _ => True
mask4 := fun _ _ _ _ _ => True
writeMask := fun _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
attention_llamaattention_llama_fwd_closed_form_correctMathematical execution | Configured model / stage Separate causal and non-causal forward value contracts. Each fixes a masking mode and context/block relations; the Realizes surface is a value contract without the KernelIO frame. | Statement and preconditionsspecification attention_llama_fwd_closed_form_correct
(Q K V Out : RegionName) (s : BlockState) (sm_scale : ℝ)
(stride_qz stride_qh stride_qm stride_qk
stride_kz stride_kh stride_kn stride_kk
stride_vz stride_vh stride_vk stride_vn
stride_oz stride_oh stride_om stride_on
N_HEAD H N_CTX start_position BLOCK_M BLOCK_N BLOCK_DMODEL numKVBlocks : Nat)
(hBM : 0 < BLOCK_M) (hBN : 0 < BLOCK_N)
(hSEQ : N_CTX = BLOCK_N * numKVBlocks) (hnum : 0 < numKVBlocks)
(houtinj : Function.Injective (fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(s.pids 0 * BLOCK_M + idx.1.val) * stride_om + idx.2.1.val * stride_on))
(hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (attention_llama_fwd_surface Q K V Out sm_scale
stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_vz stride_vh stride_vk stride_vn stride_oz stride_oh stride_om stride_on
N_HEAD H N_CTX start_position BLOCK_M BLOCK_N BLOCK_DMODEL).toAlgorithm?
= Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_llama_fwd_surface Q K V Out sm_scale
stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_vz stride_vh stride_vk stride_vn stride_oz stride_oh stride_om stride_on
N_HEAD H N_CTX start_position BLOCK_M BLOCK_N BLOCK_DMODEL)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => alActive s H BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(Out, alOutOffset s N_HEAD stride_oz stride_oh stride_om stride_on BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
attentionReal
(alQTileG s Q N_HEAD stride_qz stride_qh stride_qm stride_qk H BLOCK_M BLOCK_DMODEL)
(alKTileG s K N_HEAD stride_kz stride_kh stride_kn stride_kk N_CTX BLOCK_DMODEL)
(alVTileG s V N_HEAD stride_vz stride_vh stride_vk stride_vn N_CTX BLOCK_DMODEL)
sm_scale idx)
Referenced model/specification definitions: Python source |
|---|
attention_llamaattention_llama_fwd_causal_closed_form_correctMathematical execution | Configured model / stage Separate causal and non-causal forward value contracts. Each fixes a masking mode and context/block relations; the Realizes surface is a value contract without the KernelIO frame. | Statement and preconditionsspecification attention_llama_fwd_causal_closed_form_correct
(Q K V Out : RegionName) (s : BlockState) (sm_scale : ℝ)
(stride_qz stride_qh stride_qm stride_qk
stride_kz stride_kh stride_kn stride_kk
stride_vz stride_vh stride_vk stride_vn
stride_oz stride_oh stride_om stride_on
N_HEAD H N_CTX start_position BLOCK_M BLOCK_N BLOCK_DMODEL numCausalBlocks : Nat)
(hBM : 0 < BLOCK_M) (hBN : 0 < BLOCK_N)
(hspanEq : (s.pids 0 + 1) * BLOCK_N + start_position = BLOCK_N * numCausalBlocks)
(hspanle : BLOCK_N * numCausalBlocks ≤ N_CTX)
(hsp : start_position ≤ s.pids 0 * BLOCK_M)
(houtinj : Function.Injective (fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(s.pids 0 * BLOCK_M + idx.1.val) * stride_om + idx.2.1.val * stride_on))
(hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (attention_llama_fwd_causal_surface Q K V Out sm_scale
stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_vz stride_vh stride_vk stride_vn stride_oz stride_oh stride_om stride_on
N_HEAD H N_CTX start_position BLOCK_M BLOCK_N BLOCK_DMODEL).toAlgorithm?
= Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_llama_fwd_causal_surface Q K V Out sm_scale
stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_vz stride_vh stride_vk stride_vn stride_oz stride_oh stride_om stride_on
N_HEAD H N_CTX start_position BLOCK_M BLOCK_N BLOCK_DMODEL)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => alActive s H BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(Out, alOutOffset s N_HEAD stride_oz stride_oh stride_om stride_on BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
attentionRealCausalBlock (s.pids 0 * BLOCK_M - start_position)
(alQTileG s Q N_HEAD stride_qz stride_qh stride_qm stride_qk H BLOCK_M BLOCK_DMODEL)
(alKTileG s K N_HEAD stride_kz stride_kh stride_kn stride_kk
(BLOCK_N * numCausalBlocks) BLOCK_DMODEL)
(alVTileG s V N_HEAD stride_vz stride_vh stride_vk stride_vn
(BLOCK_N * numCausalBlocks) BLOCK_DMODEL)
sm_scale idx)
Referenced model/specification definitions: Python source |
|---|
attention_scoreattention_score_python_case1_output_summary_generalMathematical execution | Configured model / stage Case-1 score configuration only: the displayed flag tuple and BM=BN specialization. Other Python flag combinations are not covered by these headlines. | Statement and preconditionsspecification attention_score_python_case1_output_summary_general
(Q K M Out : RegionName)
(stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_oz stride_oh stride_on Z H H_KV N_CTX ROUND_CTX NKV_CTX swo sws
BN BD : Nat) (sm_scale : ℝ)
(hBNpos : 0 < BN) (hdvd : BN ∣ ROUND_CTX)
(s : BlockState) (hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (attention_score_kernel Q K M Out
stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_oz stride_oh stride_on Z H H_KV N_CTX ROUND_CTX NKV_CTX swo sws
BN BD BN sm_scale Bool.true Bool.false Bool.true Bool.true rfl).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attention_score_kernel Q K M Out
stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_oz stride_oh stride_on Z H H_KV N_CTX ROUND_CTX NKV_CTX swo sws
BN BD BN sm_scale Bool.true Bool.false Bool.true Bool.true rfl)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BN => case1OutActiveG s BN NKV_CTX i)
(fun i => (Out, case1OutStoreOffsetG s H BN stride_oz stride_oh i)))
(expected := fun i : Fin BN =>
case1OutClosedFormG s Q K M sm_scale H H_KV ROUND_CTX BN BN BD
stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kk stride_kn swo sws i)
Referenced model/specification definitions: Python source |
|---|
attention_scoreattention_score_case1_io_correctnessAbstract cast/store rounding | Configured model / stage Case-1 score configuration only: the displayed flag tuple and BM=BN specialization. Other Python flag combinations are not covered by these headlines. | Statement and preconditionsspecification attention_score_case1_io_correctness (R : RoundingModel)
(Q K M Out : RegionName)
(stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_oz stride_oh stride_on Z H H_KV N_CTX ROUND_CTX NKV_CTX swo sws
BN BD : Nat) (sm_scale : ℝ)
(hBN : 0 < BN) (hRC : 0 < ROUND_CTX) (hdvd : BN ∣ ROUND_CTX) :
attentionScoreCase1IO Q K M Out
stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_oz stride_oh stride_on Z H H_KV N_CTX ROUND_CTX NKV_CTX swo sws BN BD sm_scale ⊨[R]
fun p₀ _ _ xs ys zs j =>
attentionScoreCase1IOSpec BN BD (ROUND_CTX / BN) swo sws
(Nat.div_pos (Nat.le_of_dvd hRC hdvd) hBN) sm_scale p₀ xs ys zs j
IO definition: attentionScoreCase1IOdef attentionScoreCase1IO (Q K M Out : RegionName)
(stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_oz stride_oh stride_on Z H H_KV N_CTX ROUND_CTX NKV_CTX swo sws
BN BD : Nat) (sm_scale : ℝ) : StreamMasked3DKernelIO₃ where
kernel := attention_score_kernel Q K M Out
stride_qz stride_qh stride_qm stride_qk stride_kz stride_kh stride_kn stride_kk
stride_oz stride_oh stride_on Z H H_KV N_CTX ROUND_CTX NKV_CTX swo sws
BN BD BN sm_scale Bool.true Bool.false Bool.true Bool.true rfl
inp1 := Q
inp2 := K
inp3 := M
out := Out
T := ROUND_CTX / BN
B1 := BN * BD
B2 := BD * BN
B3 := BN
C := BN
read1 := fun _ p₁ _ t j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (t.val * BN + j.val / BD) * stride_qm + (j.val % BD) * stride_qk
read2 := fun p₀ p₁ _ _ j =>
p₁ / H * stride_kz + p₁ % H / (H / H_KV) * stride_kh
+ (j.val / BN) * stride_kk + (p₀ * BN + j.val % BN) * stride_kn
read3 := fun _ p₁ _ t j => p₁ * ROUND_CTX + (t.val * BN + j.val)
write := fun p₀ p₁ _ j => p₁ / H * stride_oz + p₁ % H * stride_oh + (p₀ * BN + j.val)
mask1 := fun _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ => True
mask3 := fun _ _ _ _ _ => True
writeMask := fun p₀ _ _ j => p₀ * BN + j.val < NKV_CTX
Referenced model/specification definitions: Python source |
|---|
attn_fwd_causalattn_fwd_causal_output_summary_generalMathematical execution | Configured model / stage Causal forward attention with the stated context/block relations and contiguous head strides. The IO theorem requires identity fp16 rounding. | Statement and preconditionsspecification attn_fwd_causal_output_summary_general
(Q K V QScale KScale Out : RegionName) (s : BlockState)
(stride_qz stride_qh H HEAD_DIM N_CTX BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE Z numKVBlocks : Nat)
(hBD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M)
(hN : N_CTX = BLOCK_N * numKVBlocks) (hnum : 0 < numKVBlocks)
(hOutInj : Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
outOffset s H stride_qz stride_qh HEAD_DIM 1 BLOCK_M idx))
(hundef : ∀ rg o, s.undef rg o = 0)
(hsb : afcScoreBoundG
(qTileAFCmG s Q stride_qz stride_qh H HEAD_DIM N_CTX BLOCK_M BLOCK_DMODEL HEAD_ACTIVE)
(kTileAFCG s K stride_qz stride_qh H HEAD_DIM (BLOCK_N * numKVBlocks) BLOCK_DMODEL)
(vTileAFCmG s V stride_qz stride_qh H HEAD_DIM (BLOCK_N * numKVBlocks) BLOCK_DMODEL HEAD_ACTIVE)
(keyScaleAFCG s QScale KScale N_CTX BLOCK_M BLOCK_N numKVBlocks) (qStartAFCG s BLOCK_M)) :
(∃ alg, (attn_fwd_causal_surface Q K V QScale KScale Out
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
Z H N_CTX HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attn_fwd_causal_surface Q K V QScale KScale Out
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
Z H N_CTX HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => active s N_CTX HEAD_ACTIVE BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => (Out,
outOffset s H stride_qz stride_qh HEAD_DIM 1 BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
attnFwdCausalOutSpecG s Q K V stride_qz stride_qh H HEAD_DIM N_CTX BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE numKVBlocks (keyScaleAFCG s QScale KScale N_CTX BLOCK_M BLOCK_N numKVBlocks) idx)
Referenced model/specification definitions: Python source |
|---|
attn_fwd_causalattn_fwd_causal_io_correctnessAbstract cast/store rounding | Configured model / stage Causal forward attention with the stated context/block relations and contiguous head strides. The IO theorem requires identity fp16 rounding. | Statement and preconditionsspecification attn_fwd_causal_io_correctness (R : RoundingModel)
(hfp16 : R.round .fp16 = id)
(Q K V QScale KScale Out : RegionName)
(stride_qz stride_qh Z H N_CTX HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE numKVBlocks : Nat)
(hBD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M)
(hN : N_CTX = BLOCK_N * numKVBlocks) (hnum : 0 < numKVBlocks)
(hBDHD : BLOCK_DMODEL ≤ HEAD_DIM)
(hsb : ∀ (p₀ : Nat) (xs : Fin numKVBlocks → Fin (BLOCK_M * BLOCK_DMODEL) → ℝ)
(ys : Fin numKVBlocks → Fin (BLOCK_DMODEL * BLOCK_N) → ℝ)
(zs : Fin numKVBlocks → Fin (BLOCK_N * BLOCK_DMODEL) → ℝ)
(ws vs : Fin numKVBlocks → Fin 1 → ℝ),
afcIOScoreBound p₀ N_CTX BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE numKVBlocks xs ys zs ws vs) :
attnFwdCausalIO Q K V QScale KScale Out stride_qz stride_qh Z H N_CTX HEAD_DIM
BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE numKVBlocks ⊨[R]
fun p₀ _ _ xs ys zs ws vs j =>
attnFwdCausalIOOutSpec p₀ N_CTX BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE numKVBlocks
xs ys zs ws vs (Lane2D.decode j)
IO definition: attnFwdCausalIOdef attnFwdCausalIO (Q K V QScale KScale Out : RegionName)
(stride_qz stride_qh Z H N_CTX HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE
numKVBlocks : Nat) : StreamMasked3DKernelIO₅ where
kernel := attn_fwd_causal_surface Q K V QScale KScale Out
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
Z H N_CTX HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE
inp1 := Q
inp2 := K
inp3 := V
inp4 := QScale
inp5 := KScale
out := Out
T := numKVBlocks
B1 := BLOCK_M * BLOCK_DMODEL
B2 := BLOCK_DMODEL * BLOCK_N
B3 := BLOCK_N * BLOCK_DMODEL
B4 := 1
B5 := 1
C := BLOCK_M * BLOCK_DMODEL
outDType := .real
read1 := fun p₀ p₁ _ _ j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (p₀ * BLOCK_M + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
read2 := fun _ p₁ _ t j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ j.val / BLOCK_N + (t.val * BLOCK_N + j.val % BLOCK_N) * HEAD_DIM
read3 := fun _ p₁ _ t j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (t.val * BLOCK_N + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
read4 := fun p₀ p₁ _ _ _ => p₁ * ((N_CTX + BLOCK_M - 1) / BLOCK_M) + p₀
read5 := fun _ p₁ _ t _ => p₁ * ((N_CTX + BLOCK_N - 1) / BLOCK_N) + t.val
write := fun p₀ p₁ _ j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (p₀ * BLOCK_M + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
mask1 := fun p₀ _ _ _ j =>
p₀ * BLOCK_M + j.val / BLOCK_DMODEL < N_CTX ∧ j.val % BLOCK_DMODEL < HEAD_ACTIVE
mask2 := fun _ _ _ t j =>
j.val % BLOCK_N < N_CTX - t.val * BLOCK_N ∧ j.val / BLOCK_N < HEAD_ACTIVE
mask3 := fun _ _ _ t j =>
j.val / BLOCK_DMODEL < N_CTX - t.val * BLOCK_N ∧ j.val % BLOCK_DMODEL < HEAD_ACTIVE
mask4 := fun _ _ _ _ _ => True
mask5 := fun _ _ _ _ _ => True
writeMask := fun p₀ _ _ j =>
p₀ * BLOCK_M + j.val / BLOCK_DMODEL < N_CTX ∧ j.val % BLOCK_DMODEL < HEAD_ACTIVE
Referenced model/specification definitions: Python source |
|---|
attn_fwd_tritonattn_fwd_triton_output_summary_generalMathematical execution | Configured model / stage Forward attention at the documented stage, layout, and context/block specialization. The IO theorem requires identity fp16 rounding. | Statement and preconditionsspecification attn_fwd_triton_output_summary_general
(Q K V QScale KScale Out : RegionName) (s : BlockState)
(stride_qz stride_qh H HEAD_DIM N_CTX BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE Z numKVBlocks : Nat)
(hBD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M)
(hN : N_CTX = BLOCK_N * numKVBlocks) (hnum : 0 < numKVBlocks)
(hOutInj : Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
outOffset s H stride_qz stride_qh HEAD_DIM 1 BLOCK_M idx))
(hundef : ∀ rg o, s.undef rg o = 0)
(hsb : aftgScoreBoundG
(qTileAFT2mG s Q stride_qz stride_qh H HEAD_DIM N_CTX BLOCK_M BLOCK_DMODEL HEAD_ACTIVE)
(kTileAFT2G s K stride_qz stride_qh H HEAD_DIM (BLOCK_N * numKVBlocks) BLOCK_DMODEL)
(vTileAFT2mG s V stride_qz stride_qh H HEAD_DIM (BLOCK_N * numKVBlocks) BLOCK_DMODEL HEAD_ACTIVE)
(keyScaleAFT2G s QScale KScale N_CTX BLOCK_M BLOCK_N numKVBlocks) (qStartAFT2G s BLOCK_M)) :
(∃ alg, (attn_fwd_triton_surface Q K V QScale KScale Out
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
Z H N_CTX HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := attn_fwd_triton_surface Q K V QScale KScale Out
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
Z H N_CTX HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => active s N_CTX HEAD_ACTIVE BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => (Out,
outOffset s H stride_qz stride_qh HEAD_DIM 1 BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
attnFwdTritonOutSpecG s Q K V stride_qz stride_qh H HEAD_DIM N_CTX BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE numKVBlocks (keyScaleAFT2G s QScale KScale N_CTX BLOCK_M BLOCK_N numKVBlocks) idx)
Referenced model/specification definitions: Python source |
|---|
attn_fwd_tritonattn_fwd_triton_io_correctnessAbstract cast/store rounding | Configured model / stage Forward attention at the documented stage, layout, and context/block specialization. The IO theorem requires identity fp16 rounding. | Statement and preconditionsspecification attn_fwd_triton_io_correctness (R : RoundingModel)
(hfp16 : R.round .fp16 = id)
(Q K V QScale KScale Out : RegionName)
(stride_qz stride_qh Z H N_CTX HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE numKVBlocks : Nat)
(hBD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M)
(hN : N_CTX = BLOCK_N * numKVBlocks) (hnum : 0 < numKVBlocks)
(hBDHD : BLOCK_DMODEL ≤ HEAD_DIM)
(hsb : ∀ (p₀ : Nat) (xs : Fin numKVBlocks → Fin (BLOCK_M * BLOCK_DMODEL) → ℝ)
(ys : Fin numKVBlocks → Fin (BLOCK_DMODEL * BLOCK_N) → ℝ)
(zs : Fin numKVBlocks → Fin (BLOCK_N * BLOCK_DMODEL) → ℝ)
(ws vs : Fin numKVBlocks → Fin 1 → ℝ),
aftgIOScoreBound p₀ N_CTX BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE numKVBlocks xs ys zs ws vs) :
attnFwdTritonIO Q K V QScale KScale Out stride_qz stride_qh Z H N_CTX HEAD_DIM
BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE numKVBlocks ⊨[R]
fun p₀ _ _ xs ys zs ws vs j =>
attnFwdTritonIOOutSpec p₀ N_CTX BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE numKVBlocks
xs ys zs ws vs (Lane2D.decode j)
IO definition: attnFwdTritonIOdef attnFwdTritonIO (Q K V QScale KScale Out : RegionName)
(stride_qz stride_qh Z H N_CTX HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE
numKVBlocks : Nat) : StreamMasked3DKernelIO₅ where
kernel := attn_fwd_triton_surface Q K V QScale KScale Out
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
stride_qz stride_qh HEAD_DIM 1 stride_qz stride_qh HEAD_DIM 1
Z H N_CTX HEAD_DIM BLOCK_M BLOCK_N BLOCK_DMODEL HEAD_ACTIVE STAGE
inp1 := Q
inp2 := K
inp3 := V
inp4 := QScale
inp5 := KScale
out := Out
T := numKVBlocks
B1 := BLOCK_M * BLOCK_DMODEL
B2 := BLOCK_DMODEL * BLOCK_N
B3 := BLOCK_N * BLOCK_DMODEL
B4 := 1
B5 := 1
C := BLOCK_M * BLOCK_DMODEL
outDType := .real
read1 := fun p₀ p₁ _ _ j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (p₀ * BLOCK_M + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
read2 := fun _ p₁ _ t j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ j.val / BLOCK_N + (t.val * BLOCK_N + j.val % BLOCK_N) * HEAD_DIM
read3 := fun _ p₁ _ t j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (t.val * BLOCK_N + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
read4 := fun p₀ p₁ _ _ _ => p₁ * ((N_CTX + BLOCK_M - 1) / BLOCK_M) + p₀
read5 := fun _ p₁ _ t _ => p₁ * ((N_CTX + BLOCK_N - 1) / BLOCK_N) + t.val
write := fun p₀ p₁ _ j =>
p₁ / H * stride_qz + p₁ % H * stride_qh
+ (p₀ * BLOCK_M + j.val / BLOCK_DMODEL) * HEAD_DIM + j.val % BLOCK_DMODEL
mask1 := fun p₀ _ _ _ j =>
p₀ * BLOCK_M + j.val / BLOCK_DMODEL < N_CTX ∧ j.val % BLOCK_DMODEL < HEAD_ACTIVE
mask2 := fun _ _ _ t j =>
j.val % BLOCK_N < N_CTX - t.val * BLOCK_N ∧ j.val / BLOCK_N < HEAD_ACTIVE
mask3 := fun _ _ _ t j =>
j.val / BLOCK_DMODEL < N_CTX - t.val * BLOCK_N ∧ j.val % BLOCK_DMODEL < HEAD_ACTIVE
mask4 := fun _ _ _ _ _ => True
mask5 := fun _ _ _ _ _ => True
writeMask := fun p₀ _ _ j =>
p₀ * BLOCK_M + j.val / BLOCK_DMODEL < N_CTX ∧ j.val % BLOCK_DMODEL < HEAD_ACTIVE
Referenced model/specification definitions: Python source |
|---|
batched_vecmat_multbatched_vecmat_closed_form_correctMathematical execution | Configured model / stage Batched vector/matrix program with reduction length BLOCK_K*numKBlocks and the stated block/dimension bounds; host launch and general tail choices are external. | Statement and preconditionsspecification batched_vecmat_closed_form_correct
(A B output : RegionName) (s : BlockState)
(_dim_m dim_n BLOCK_M BLOCK_N BLOCK_K numKBlocks : Nat)
(hBM : 0 < BLOCK_M) (hBN : 0 < BLOCK_N) (hBK : 0 < BLOCK_K)
(hInj : Function.Injective (vecmatOutOffset s dim_n BLOCK_M BLOCK_N))
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := batched_vecmat_surface A B output _dim_m dim_n (BLOCK_K * numKBlocks)
BLOCK_M BLOCK_N BLOCK_K)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
some (output, vecmatOutOffset s dim_n BLOCK_M BLOCK_N idx))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
gemvSpec s A B dim_n (BLOCK_K * numKBlocks) BLOCK_M BLOCK_N BLOCK_K numKBlocks
idx.1 idx.2.1)
Referenced model/specification definitions: Python source |
|---|
batched_vecmat_multbatched_vecmat_io_correctnessAbstract cast/store rounding | Configured model / stage Batched vector/matrix program with reduction length BLOCK_K*numKBlocks and the stated block/dimension bounds; host launch and general tail choices are external. | Statement and preconditionsspecification batched_vecmat_io_correctness (R : RoundingModel)
(A B output : RegionName)
(_dim_m dim_n BLOCK_M BLOCK_N BLOCK_K numKBlocks : Nat)
(hBK : 0 < BLOCK_K) (hBN : BLOCK_N ≤ dim_n) :
batchedVecmatIO A B output _dim_m dim_n BLOCK_M BLOCK_N BLOCK_K numKBlocks
⊨[R] fun _ _ xs ys l =>
∑ t : Fin numKBlocks, ∑ e : Fin BLOCK_K,
xs t (aLane BLOCK_M BLOCK_N BLOCK_K l e) * ys t (bLane BLOCK_M BLOCK_N BLOCK_K l e)
IO definition: batchedVecmatIOdef batchedVecmatIO (A B output : RegionName)
(_dim_m dim_n BLOCK_M BLOCK_N BLOCK_K numKBlocks : Nat) :
StreamMasked2DKernelIO₂ where
kernel := batched_vecmat_surface A B output _dim_m dim_n (BLOCK_K * numKBlocks)
BLOCK_M BLOCK_N BLOCK_K
inp1 := A
inp2 := B
out := output
T := numKBlocks
B1 := BLOCK_M * BLOCK_K
B2 := BLOCK_N * BLOCK_M * BLOCK_K
C := BLOCK_M * BLOCK_N
read1 := fun p₀ _ t l =>
(p₀ * BLOCK_M + l.val / BLOCK_K) * (BLOCK_K * numKBlocks)
+ (t.val * BLOCK_K + l.val % BLOCK_K)
read2 := fun p₀ p₁ t l =>
(p₀ * BLOCK_M + l.val / BLOCK_K % BLOCK_M) * dim_n * (BLOCK_K * numKBlocks)
+ (p₁ * BLOCK_N + l.val / BLOCK_K / BLOCK_M) * (BLOCK_K * numKBlocks)
+ (t.val * BLOCK_K + l.val % BLOCK_K)
write := fun p₀ p₁ l =>
(p₀ * BLOCK_M + l.val / BLOCK_N) * dim_n + (p₁ * BLOCK_N + l.val % BLOCK_N)
mask1 := fun _ _ _ _ => True
mask2 := fun _ _ _ _ => True
writeMask := fun _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
bgmv_expand_slicebgmv_full_output_summaryMathematical execution | Configured model / stage One selected LoRA expansion slice with explicit slice offset and tile sizes. The contract's loaded LoRA index and branch/layout assumptions must hold; it is not a host adapter-selection proof. | Statement and preconditionsspecification bgmv_full_output_summary
(input_ptr lora_ptr out_ptr : RegionName) (lora_indices : Region .int)
(li K split_n_length xm_stride xk_stride l0_stride lora_k_stride
lora_n_stride cm_stride cn_stride slice_offset BLOCK_N BLOCK_K : Nat)
(s : BlockState) (hBN : 0 < BLOCK_N)
(hoi : out_ptr ≠ input_ptr) (hol : out_ptr ≠ lora_ptr)
(hcn : 0 < cn_stride)
(hlx : s.readMemValue .int (Region.cast lora_indices) (s.pids 1) = Int.ofNat li) :
(∃ alg, (bgmv_full input_ptr lora_ptr out_ptr lora_indices K split_n_length xm_stride xk_stride
l0_stride lora_k_stride lora_n_stride cm_stride cn_stride slice_offset BLOCK_N BLOCK_K).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := bgmv_full input_ptr lora_ptr out_ptr lora_indices K split_n_length xm_stride xk_stride
l0_stride lora_k_stride lora_n_stride cm_stride cn_stride slice_offset BLOCK_N BLOCK_K)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun _ : Fin split_n_length => True)
(fun g => (out_ptr, cOff s split_n_length cm_stride cn_stride slice_offset g.val)))
(expected := fun g : Fin split_n_length =>
bgmvFullSpec s input_ptr lora_ptr li K split_n_length xm_stride xk_stride l0_stride lora_k_stride lora_n_stride BLOCK_K g.val)
Referenced model/specification definitions: Python source |
|---|
bgmv_expand_slicebgmv_expand_slice_io_correctnessAbstract cast/store rounding | Configured model / stage One selected LoRA expansion slice with explicit slice offset and tile sizes. The contract's loaded LoRA index and branch/layout assumptions must hold; it is not a host adapter-selection proof. | Statement and preconditionsspecification bgmv_expand_slice_io_correctness (R : RoundingModel)
(input_ptr lora_ptr out_ptr : RegionName) (lora_indices : Region .int)
(K split_n_length xm_stride xk_stride l0_stride lora_k_stride
lora_n_stride cm_stride cn_stride slice_offset BLOCK_N BLOCK_K : Nat)
(hBN : 0 < BLOCK_N) (hsn : 0 < split_n_length)
(hoi : out_ptr ≠ input_ptr) (hol : out_ptr ≠ lora_ptr)
(hcn : 0 < cn_stride) :
bgmvExpandSliceIO input_ptr lora_ptr out_ptr lora_indices K
split_n_length xm_stride xk_stride l0_stride lora_k_stride
lora_n_stride cm_stride cn_stride slice_offset BLOCK_N BLOCK_K
⊨[R] fun _ _ _ m xs ys t j =>
∑ k : Fin BLOCK_K,
(if k.val < K then xs t k else 0)
* (if t.val * BLOCK_N + j.val < split_n_length ∧ k.val < K
then ys t (bTileLane BLOCK_N BLOCK_K j k) else 0)
IO definition: bgmvExpandSliceIOdef bgmvExpandSliceIO (input_ptr lora_ptr out_ptr : RegionName)
(lora_indices : Region .int)
(K split_n_length xm_stride xk_stride l0_stride lora_k_stride
lora_n_stride cm_stride cn_stride slice_offset BLOCK_N BLOCK_K : Nat) :
StreamMetaEmitMasked3DKernelIO₂ where
kernel := bgmv_full input_ptr lora_ptr out_ptr lora_indices K
split_n_length xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride slice_offset BLOCK_N BLOCK_K
inp1 := input_ptr
inp2 := lora_ptr
out := out_ptr
nMeta := 1
sty := fun _ => .int
mbuf := fun _ => Region.cast lora_indices
mwin := fun _ _ pid₁ _ => pid₁
T := bgmvNumSteps split_n_length BLOCK_N
B1 := BLOCK_K
B2 := BLOCK_N * BLOCK_K
C := BLOCK_N
read1 := fun _ pid₁ _ _ _ k => pid₁ * xm_stride + k.val * xk_stride
read2 := fun pid₀ _ _ m t l =>
l0_stride * (m ⟨0, Nat.one_pos⟩).toNat
+ pid₀ * split_n_length * lora_k_stride
+ (t.val * BLOCK_N + l.val / BLOCK_K) * lora_k_stride
+ l.val % BLOCK_K * lora_n_stride
write := fun pid₀ pid₁ _ _ t j =>
pid₁ * cm_stride + pid₀ * split_n_length + slice_offset * cn_stride
+ (t.val * BLOCK_N + j.val) * cn_stride
mask1 := fun _ _ _ _ _ k => k.val < K
mask2 := fun _ _ _ _ t l =>
t.val * BLOCK_N + l.val / BLOCK_K < split_n_length ∧ l.val % BLOCK_K < K
writeMask := fun _ _ _ m t j =>
m ⟨0, Nat.one_pos⟩ ≠ -1 ∧ t.val * BLOCK_N + j.val < split_n_length
Referenced model/specification definitions: Python source |
|---|
bgmv_shrink_kernelbgmv_shrink_kernel_output_summary_generalMathematical execution | Configured model / stage Shrink output and store-branch contracts with explicit SPLIT_K addressing. Atomic projection/readback facts do not establish concurrent multi-program atomic execution. | Statement and preconditionsspecification bgmv_shrink_kernel_output_summary_general
(input_ptr lora_ptr out_ptr : RegionName)
(N K : Nat) (lora_indices_int : Region .int) (lora_indices : Region .nat)
(scaling : ℝ)
(xm_stride l0_stride lora_k_stride lora_n_stride cm_stride cn_stride
BLOCK_N BLOCK_K SPLIT_K : Nat) (SPLIT_K_ONE : Bool)
(s : BlockState)
(hBK : 0 < BLOCK_K) (hSK : 0 < SPLIT_K) (hcn : 0 < cn_stride) :
-- (1) the faithful guarded surface (sentinel guard + both constexpr tail
-- branches) lowers to the algorithm layer
(∃ alg, (bgmv_shrink_surface input_ptr lora_ptr out_ptr N K
lora_indices_int scaling xm_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K SPLIT_K SPLIT_K_ONE).toAlgorithm?
= Except.ok alg) ∧
-- (2) both per-branch proof surfaces lower to the algorithm layer
(∃ alg, (bgmv_shrink_store_surface input_ptr lora_ptr out_ptr N K
lora_indices scaling xm_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K SPLIT_K).toAlgorithm?
= Except.ok alg) ∧
(∃ alg, (bgmv_shrink_atomic_surface input_ptr lora_ptr out_ptr N K
lora_indices scaling xm_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K SPLIT_K).toAlgorithm?
= Except.ok alg) ∧
-- (3) the sentinel early-return path writes nothing
(∀ s', s.readMemValue .int lora_indices_int (s.pids 1) = (-1 : Int) →
exec (bgmv_shrink_surface input_ptr lora_ptr out_ptr N K lora_indices_int
scaling xm_stride l0_stride lora_k_stride lora_n_stride cm_stride
cn_stride BLOCK_N BLOCK_K SPLIT_K SPLIT_K_ONE) s = some s' →
s'.mem = s.mem) ∧
-- (4) SPLIT_K = 1: masked store of the genuine contraction
ComputeCorrect.Realizes_without_Rounding
(kernel := bgmv_shrink_store_surface input_ptr lora_ptr out_ptr N K
lora_indices scaling xm_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K SPLIT_K)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun n : Fin BLOCK_N => n.val < N)
(fun n => (out_ptr, outOff s cm_stride cn_stride n.val)))
(expected := fun n : Fin BLOCK_N =>
shrinkSpec s input_ptr lora_ptr lora_indices scaling K xm_stride
l0_stride lora_k_stride lora_n_stride BLOCK_K SPLIT_K n.val) ∧
-- (5) SPLIT_K > 1: masked atomic add of the genuine contraction
ComputeCorrect.Realizes_without_Rounding
(kernel := bgmv_shrink_atomic_surface input_ptr lora_ptr out_ptr N K
lora_indices scaling xm_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K SPLIT_K)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun n : Fin BLOCK_N => n.val < N)
(fun n => (out_ptr, outOff s cm_stride cn_stride n.val)))
(expected := fun n : Fin BLOCK_N =>
s.readMem out_ptr (outOff s cm_stride cn_stride n.val)
+ shrinkSpec s input_ptr lora_ptr lora_indices scaling K xm_stride
l0_stride lora_k_stride lora_n_stride BLOCK_K SPLIT_K n.val)
Referenced model/specification definitions: Python source |
|---|
bgmv_shrink_kernelbgmv_shrink_store_io_correctnessAbstract cast/store rounding | Configured model / stage Shrink output and store-branch contracts with explicit SPLIT_K addressing. Atomic projection/readback facts do not establish concurrent multi-program atomic execution. | Statement and preconditionsspecification bgmv_shrink_store_io_correctness (R : RoundingModel)
(input_ptr lora_ptr out_ptr : RegionName)
(N K : Nat) (lora_indices : Region .nat) (scaling : ℝ)
(xm_stride l0_stride lora_k_stride lora_n_stride cm_stride cn_stride
BLOCK_N BLOCK_K SPLIT_K : Nat)
(hS : 0 < BLOCK_K * SPLIT_K) (hcn : 0 < cn_stride) :
bgmvShrinkStoreIO input_ptr lora_ptr out_ptr N K lora_indices scaling
xm_stride l0_stride lora_k_stride lora_n_stride cm_stride cn_stride
BLOCK_N BLOCK_K SPLIT_K
⊨[R] fun pid₀ _ _ _ xs ys j =>
scaling * ∑ t : Fin (numKIters K (BLOCK_K * SPLIT_K)),
∑ e : Fin BLOCK_K,
if t.val * (BLOCK_K * SPLIT_K) + (e.val + pid₀ * BLOCK_K) < K
then xs t e * ys t (bStreamLane BLOCK_N BLOCK_K j e)
else 0
IO definition: bgmvShrinkStoreIOdef bgmvShrinkStoreIO (input_ptr lora_ptr out_ptr : RegionName)
(N K : Nat) (lora_indices : Region .nat) (scaling : ℝ)
(xm_stride l0_stride lora_k_stride lora_n_stride cm_stride cn_stride
BLOCK_N BLOCK_K SPLIT_K : Nat) : StreamMetaMasked3DKernelIO₂ where
kernel := bgmv_shrink_store_surface input_ptr lora_ptr out_ptr N K
lora_indices scaling xm_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K SPLIT_K
inp1 := input_ptr
inp2 := lora_ptr
out := out_ptr
nMeta := 1
sty := fun _ => .nat
mbuf := fun _ => Region.cast lora_indices
mwin := fun _ _ pid₁ _ => pid₁
T := numKIters K (BLOCK_K * SPLIT_K)
B1 := BLOCK_K
B2 := BLOCK_N * BLOCK_K
C := BLOCK_N
outDType := .real
read1 := fun pid₀ pid₁ _ _ t e =>
pid₁ * xm_stride + (t.val * (BLOCK_K * SPLIT_K) + (e.val + pid₀ * BLOCK_K))
read2 := fun pid₀ _ _ m t l =>
l0_stride * m 0 + l.val / BLOCK_K * lora_k_stride
+ (t.val * (BLOCK_K * SPLIT_K) + (l.val % BLOCK_K + pid₀ * BLOCK_K))
* lora_n_stride
write := fun _ pid₁ _ _ j => pid₁ * cm_stride + j.val * cn_stride
mask1 := fun pid₀ _ _ _ t e =>
t.val * (BLOCK_K * SPLIT_K) + (e.val + pid₀ * BLOCK_K) < K
mask2 := fun pid₀ _ _ _ t l =>
l.val / BLOCK_K < N ∧
t.val * (BLOCK_K * SPLIT_K) + (l.val % BLOCK_K + pid₀ * BLOCK_K) < K
writeMask := fun _ _ _ _ j => j.val < N
Referenced model/specification definitions: Python source |
|---|
block_sparse_attnblock_sparse_attn_output_closed_form_summary_generalMathematical execution | Configured model / stage Selected sparse-layout forward value contract, plus separate output-store slices. Sparse metadata/layout premises and selected dimensions remain part of the statement. | Statement and preconditionsspecification block_sparse_attn_output_closed_form_summary_general
(Out Q K V : RegionName) (R C : Region .nat) (s : BlockState)
(BLOCK_M BLOCK_D BLOCK_N num_heads num_kv_heads num_layout total_seq_len : Nat)
(rowStrideH colStrideH sqb sqh sqm skb skh skn svb svh svn sob soh som : Nat) (scale : ℝ)
(hundef : ∀ rg o, s.undef rg o = 0)
(hBN : 0 < BLOCK_N)
(start_l end_l : Nat)
(hStartL : s.readMemValue .nat R.cast
(s.pids 1 % num_heads % num_layout * rowStrideH + s.pids 0) = start_l)
(hEndL : s.readMemValue .nat R.cast
(s.pids 1 % num_heads % num_layout * rowStrideH + s.pids 0 + 1) = end_l)
(hsle : start_l ≤ end_l) (hN : 0 < end_l - start_l)
(hinj : Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_D] => outOffset s num_heads sob soh som BLOCK_M idx))
(hdisj : ∀ a b : TileIndex [BLOCK_M, BLOCK_D],
outOffset s num_heads sob soh som BLOCK_M a
≠ out2Offset s num_heads sob soh som BLOCK_M BLOCK_D b)
(hVis0 : ∀ idx : TileIndex [BLOCK_M, BLOCK_D], active s total_seq_len BLOCK_M idx →
selKeyGlobal s C (s.pids 1 % num_heads % num_layout) num_heads start_l BLOCK_N 0
≤ s.pids 0 * BLOCK_M + idx.1.val)
(hstep : ∀ (i : Nat) (st : BlockState), start_l ≤ i → i < end_l →
bsaInvariantG Out Q K V R C BLOCK_M BLOCK_D BLOCK_N num_heads num_kv_heads num_layout
total_seq_len sqb sqh sqm skb skh skn svb svh svn (s.pids 0 * BLOCK_M) (end_l - start_l)
(fun r : Fin (BLOCK_N * (end_l - start_l)) =>
selKeyGlobal s C (s.pids 1 % num_heads % num_layout) num_heads start_l BLOCK_N r.val)
(fun jx : TileIndex [BLOCK_M, 2 * BLOCK_D] =>
qTileBSA s Q num_heads sqb sqh sqm BLOCK_M jx.1 jx.2.1.val)
(fun jx : TileIndex [BLOCK_N * (end_l - start_l), 2 * BLOCK_D] =>
kRowBSA s K num_heads num_kv_heads skb skh skn
(selKeyGlobal s C (s.pids 1 % num_heads % num_layout) num_heads start_l BLOCK_N jx.1.val) jx.2.1.val)
(fun jx : TileIndex [BLOCK_N * (end_l - start_l), BLOCK_D] =>
vRowBSA s V num_heads num_kv_heads svb svh svn
(selKeyGlobal s C (s.pids 1 % num_heads % num_layout) num_heads start_l BLOCK_N jx.1.val) (0 + jx.2.1.val))
(fun jx : TileIndex [BLOCK_N * (end_l - start_l), BLOCK_D] =>
vRowBSA s V num_heads num_kv_heads svb svh svn
(selKeyGlobal s C (s.pids 1 % num_heads % num_layout) num_heads start_l BLOCK_N jx.1.val) (BLOCK_D + jx.2.1.val))
scale s (i - start_l) st →
∃ st', stepStmts (bsaLoopBodyG C BLOCK_M BLOCK_D BLOCK_N num_heads colStrideH skn svn total_seq_len scale)
(st.setReg "col_idx_idx" .nat [] (Tile.scalar i)) = some st'
∧ bsaInvariantG Out Q K V R C BLOCK_M BLOCK_D BLOCK_N num_heads num_kv_heads num_layout
total_seq_len sqb sqh sqm skb skh skn svb svh svn (s.pids 0 * BLOCK_M) (end_l - start_l)
(fun r : Fin (BLOCK_N * (end_l - start_l)) =>
selKeyGlobal s C (s.pids 1 % num_heads % num_layout) num_heads start_l BLOCK_N r.val)
(fun jx : TileIndex [BLOCK_M, 2 * BLOCK_D] =>
qTileBSA s Q num_heads sqb sqh sqm BLOCK_M jx.1 jx.2.1.val)
(fun jx : TileIndex [BLOCK_N * (end_l - start_l), 2 * BLOCK_D] =>
kRowBSA s K num_heads num_kv_heads skb skh skn
(selKeyGlobal s C (s.pids 1 % num_heads % num_layout) num_heads start_l BLOCK_N jx.1.val) jx.2.1.val)
(fun jx : TileIndex [BLOCK_N * (end_l - start_l), BLOCK_D] =>
vRowBSA s V num_heads num_kv_heads svb svh svn
(selKeyGlobal s C (s.pids 1 % num_heads % num_layout) num_heads start_l BLOCK_N jx.1.val) (0 + jx.2.1.val))
(fun jx : TileIndex [BLOCK_N * (end_l - start_l), BLOCK_D] =>
vRowBSA s V num_heads num_kv_heads svb svh svn
(selKeyGlobal s C (s.pids 1 % num_heads % num_layout) num_heads start_l BLOCK_N jx.1.val) (BLOCK_D + jx.2.1.val))
scale s (i - start_l + 1) st') :
(ComputeCorrect.Realizes_without_Rounding
(kernel := block_sparse_attention_kernel Out Q K V R C
rowStrideH colStrideH num_layout scale sqb sqh sqm skb skh skn svb svh svn sob soh som
num_heads num_kv_heads total_seq_len BLOCK_M BLOCK_N BLOCK_D 2 Bool.true Bool.true)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_D] => active s total_seq_len BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_D] => (Out, outOffset s num_heads sob soh som BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_D] =>
blockSparseAttnClosedForm s Q K V C num_heads num_kv_heads sqb sqh sqm skb skh skn svb svh svn
(s.pids 1 % num_heads % num_layout) num_heads start_l (end_l - start_l) (2 * BLOCK_D) BLOCK_M BLOCK_N
0 scale idx.1 idx.2.1.val)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := block_sparse_attention_kernel Out Q K V R C
rowStrideH colStrideH num_layout scale sqb sqh sqm skb skh skn svb svh svn sob soh som
num_heads num_kv_heads total_seq_len BLOCK_M BLOCK_N BLOCK_D 2 Bool.true Bool.true)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_D] => active s total_seq_len BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_D] => (Out, out2Offset s num_heads sob soh som BLOCK_M BLOCK_D idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_D] =>
blockSparseAttnClosedForm s Q K V C num_heads num_kv_heads sqb sqh sqm skb skh skn svb svh svn
(s.pids 1 % num_heads % num_layout) num_heads start_l (end_l - start_l) (2 * BLOCK_D) BLOCK_M BLOCK_N
BLOCK_D scale idx.1 idx.2.1.val))
Referenced model/specification definitions: Python source |
|---|
block_sparse_attnblock_sparse_attn_output_stores_io_correctnessMathematical execution | Precomputed-input slice First and second output-store slices consume precomputed Acc/Acc2 tiles; this headline does not compute sparse attention. | Statement and preconditionsspecification block_sparse_attn_output_stores_io_correctness
(Acc Acc2 Out : RegionName)
(num_heads total_seq_len stride_acc_b stride_acc_h stride_acc_m stride_acc_d
stride_ob stride_oh stride_om BLOCK_M BLOCK_D : Nat)
(hInj1 : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_D] =>
p₁ / num_heads * stride_ob + p₁ % num_heads * stride_oh
+ (p₀ * BLOCK_M + idx.1.val) * stride_om + idx.2.1.val))
(hInj2 : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_D] =>
p₁ / num_heads * stride_ob + p₁ % num_heads * stride_oh
+ (p₀ * BLOCK_M + idx.1.val) * stride_om + BLOCK_D + idx.2.1.val)) :
(out_storeIO Acc Out num_heads total_seq_len stride_acc_b stride_acc_h
stride_acc_m stride_acc_d stride_ob stride_oh stride_om BLOCK_M BLOCK_D
⊨ fun _p₀ _p₁ xs idx => xs idx) ∧
(out2_storeIO Acc2 Out num_heads total_seq_len stride_acc_b stride_acc_h
stride_acc_m stride_acc_d stride_ob stride_oh stride_om BLOCK_M BLOCK_D
⊨ fun _p₀ _p₁ xs idx => xs idx)
IO definition: out_storeIOdef out_storeIO (Acc Out : RegionName)
(num_heads total_seq_len stride_acc_b stride_acc_h stride_acc_m stride_acc_d
stride_ob stride_oh stride_om BLOCK_M BLOCK_D : Nat) : Masked3DTileKernelIO₁ where
kernel := block_sparse_attn_output_store_slice Acc Out num_heads total_seq_len stride_acc_b stride_acc_h
stride_acc_m stride_acc_d stride_ob stride_oh stride_om BLOCK_M BLOCK_D
inp := Acc
out := Out
shape := [BLOCK_M, BLOCK_D]
read := fun p₀ p₁ _p₂ idx =>
p₁ / num_heads * stride_acc_b + p₁ % num_heads * stride_acc_h
+ (p₀ * BLOCK_M + idx.1.val) * stride_acc_m + idx.2.1.val * stride_acc_d
write := fun p₀ p₁ _p₂ idx =>
p₁ / num_heads * stride_ob + p₁ % num_heads * stride_oh
+ (p₀ * BLOCK_M + idx.1.val) * stride_om + idx.2.1.val
mask := fun p₀ _p₁ _p₂ idx => p₀ * BLOCK_M + idx.1.val < total_seq_len
IO definition: out2_storeIOdef out2_storeIO (Acc2 Out : RegionName)
(num_heads total_seq_len stride_acc_b stride_acc_h stride_acc_m stride_acc_d
stride_ob stride_oh stride_om BLOCK_M BLOCK_D : Nat) : Masked3DTileKernelIO₁ where
kernel := block_sparse_attn_output_store_second_slice Acc2 Out num_heads total_seq_len stride_acc_b stride_acc_h
stride_acc_m stride_acc_d stride_ob stride_oh stride_om BLOCK_M BLOCK_D
inp := Acc2
out := Out
shape := [BLOCK_M, BLOCK_D]
read := fun p₀ p₁ _p₂ idx =>
p₁ / num_heads * stride_acc_b + p₁ % num_heads * stride_acc_h
+ (p₀ * BLOCK_M + idx.1.val) * stride_acc_m + idx.2.1.val * stride_acc_d
write := fun p₀ p₁ _p₂ idx =>
p₁ / num_heads * stride_ob + p₁ % num_heads * stride_oh
+ (p₀ * BLOCK_M + idx.1.val) * stride_om + BLOCK_D + idx.2.1.val
mask := fun p₀ _p₁ _p₂ idx => p₀ * BLOCK_M + idx.1.val < total_seq_len
Referenced model/specification definitions: Python source |
|---|
block_sparse_attnblock_sparse_attn_output_stores_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding contracts for stores of precomputed Acc/Acc2 tiles; the attention producer is outside these two IO contracts. | Statement and preconditionsspecification block_sparse_attn_output_stores_io_correctnessR (R : RoundingModel)
(Acc Acc2 Out : RegionName)
(num_heads total_seq_len stride_acc_b stride_acc_h stride_acc_m stride_acc_d
stride_ob stride_oh stride_om BLOCK_M BLOCK_D : Nat)
(hInj1 : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_D] =>
p₁ / num_heads * stride_ob + p₁ % num_heads * stride_oh
+ (p₀ * BLOCK_M + idx.1.val) * stride_om + idx.2.1.val))
(hInj2 : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_D] =>
p₁ / num_heads * stride_ob + p₁ % num_heads * stride_oh
+ (p₀ * BLOCK_M + idx.1.val) * stride_om + BLOCK_D + idx.2.1.val)) :
(out_storeIO Acc Out num_heads total_seq_len stride_acc_b stride_acc_h
stride_acc_m stride_acc_d stride_ob stride_oh stride_om BLOCK_M BLOCK_D
⊨[R, FloatDType.real] fun _p₀ _p₁ xs idx => xs idx) ∧
(out2_storeIO Acc2 Out num_heads total_seq_len stride_acc_b stride_acc_h
stride_acc_m stride_acc_d stride_ob stride_oh stride_om BLOCK_M BLOCK_D
⊨[R, FloatDType.real] fun _p₀ _p₁ xs idx => xs idx)
IO definition: out_storeIOdef out_storeIO (Acc Out : RegionName)
(num_heads total_seq_len stride_acc_b stride_acc_h stride_acc_m stride_acc_d
stride_ob stride_oh stride_om BLOCK_M BLOCK_D : Nat) : Masked3DTileKernelIO₁ where
kernel := block_sparse_attn_output_store_slice Acc Out num_heads total_seq_len stride_acc_b stride_acc_h
stride_acc_m stride_acc_d stride_ob stride_oh stride_om BLOCK_M BLOCK_D
inp := Acc
out := Out
shape := [BLOCK_M, BLOCK_D]
read := fun p₀ p₁ _p₂ idx =>
p₁ / num_heads * stride_acc_b + p₁ % num_heads * stride_acc_h
+ (p₀ * BLOCK_M + idx.1.val) * stride_acc_m + idx.2.1.val * stride_acc_d
write := fun p₀ p₁ _p₂ idx =>
p₁ / num_heads * stride_ob + p₁ % num_heads * stride_oh
+ (p₀ * BLOCK_M + idx.1.val) * stride_om + idx.2.1.val
mask := fun p₀ _p₁ _p₂ idx => p₀ * BLOCK_M + idx.1.val < total_seq_len
IO definition: out2_storeIOdef out2_storeIO (Acc2 Out : RegionName)
(num_heads total_seq_len stride_acc_b stride_acc_h stride_acc_m stride_acc_d
stride_ob stride_oh stride_om BLOCK_M BLOCK_D : Nat) : Masked3DTileKernelIO₁ where
kernel := block_sparse_attn_output_store_second_slice Acc2 Out num_heads total_seq_len stride_acc_b stride_acc_h
stride_acc_m stride_acc_d stride_ob stride_oh stride_om BLOCK_M BLOCK_D
inp := Acc2
out := Out
shape := [BLOCK_M, BLOCK_D]
read := fun p₀ p₁ _p₂ idx =>
p₁ / num_heads * stride_acc_b + p₁ % num_heads * stride_acc_h
+ (p₀ * BLOCK_M + idx.1.val) * stride_acc_m + idx.2.1.val * stride_acc_d
write := fun p₀ p₁ _p₂ idx =>
p₁ / num_heads * stride_ob + p₁ % num_heads * stride_oh
+ (p₀ * BLOCK_M + idx.1.val) * stride_om + BLOCK_D + idx.2.1.val
mask := fun p₀ _p₁ _p₂ idx => p₀ * BLOCK_M + idx.1.val < total_seq_len
Referenced model/specification definitions: Python source |
|---|
bmm_chunk_bwdbmm_chunk_bwd_output_summary_generalMathematical execution | Configured model / stage Chunked backward matrix product with reduction extent BCS*numCSBlocks, explicit batch/group layout, masks, and injective output addresses. | Statement and preconditionsspecification bmm_chunk_bwd_output_summary_general
(A Dout Db : RegionName) (s : BlockState)
(chunk_size ngroups SAB SAS SAH SAK SDB SDC SDH SDM SDN SOB SOS SOH SOK BM BN BCS numCSBlocks K : Nat) (hBCS : 0 < BCS)
(hInj : Function.Injective (dbOffset (s.pids 1) (pidC (s.pids 2) ngroups) (pidH (s.pids 2) ngroups)
chunk_size SOB SOS SOH SOK (pidM (s.pids 0) K BN) (pidN (s.pids 0) K BN) BM BN))
(hmlt : ∀ i : Fin BM, rowIndex (pidM (s.pids 0) K BN) BM i < BCS * numCSBlocks)
(hmlt' : ∀ i : Fin BM, rowIndex (pidM (s.pids 0) K BN) BM i < chunk_size)
(hnlt : ∀ j : Fin BN, colIndex (pidN (s.pids 0) K BN) BN j < K)
(hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (bbwd_matmul_surface A Dout Db chunk_size (BCS * numCSBlocks) K ngroups
SAB SAS SAH SAK SDB SDC SDH SDM SDN SOB SOS SOH SOK BM BN BCS).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := bbwd_matmul_surface A Dout Db chunk_size (BCS * numCSBlocks) K ngroups
SAB SAS SAH SAK SDB SDC SDH SDM SDN SOB SOS SOH SOK BM BN BCS)
(initialState := s)
(write := fun idx : TileIndex [BM, BN] =>
some (Db, dbOffset (s.pids 1) (pidC (s.pids 2) ngroups) (pidH (s.pids 2) ngroups)
chunk_size SOB SOS SOH SOK (pidM (s.pids 0) K BN) (pidN (s.pids 0) K BN) BM BN idx))
(expected := fun idx : TileIndex [BM, BN] =>
dbCell s Dout A (s.pids 1) (pidC (s.pids 2) ngroups) (pidH (s.pids 2) ngroups)
(pidM (s.pids 0) K BN) (pidN (s.pids 0) K BN)
chunk_size BM BN SDB SDC SDH SDN SDM SAB SAS SAH SAK BCS numCSBlocks idx)
Referenced model/specification definitions: Python source |
|---|
bmm_chunk_bwdbmm_chunk_bwd_io_correctnessAbstract cast/store rounding | Configured model / stage Chunked backward matrix product with reduction extent BCS*numCSBlocks, explicit batch/group layout, masks, and injective output addresses. | Statement and preconditionsspecification bmm_chunk_bwd_io_correctness (R : RoundingModel)
(A Dout Db : RegionName)
(chunk_size K ngroups SAB SAS SAH SAK SDB SDC SDH SDM SDN SOB SOS SOH SOK
BM BN BCS numCSBlocks : Nat)
(hBCS : 0 < BCS)
(hOutInj : ∀ pid₀ pid₁ pid₂ : Nat,
Function.Injective (dbOffset pid₁ (pidC pid₂ ngroups) (pidH pid₂ ngroups)
chunk_size SOB SOS SOH SOK (pidM pid₀ K BN) (pidN pid₀ K BN) BM BN)) :
bmm_chunk_bwd_IO A Dout Db chunk_size K ngroups SAB SAS SAH SAK SDB SDC SDH SDM SDN
SOB SOS SOH SOK BM BN BCS numCSBlocks ⊨[R]
fun pid₀ _ _ _ xs ys l =>
bbwdStreamSum (pidM pid₀ K BN) (pidN pid₀ K BN) BM BN BCS numCSBlocks chunk_size K
xs ys l
IO definition: bmm_chunk_bwd_IOdef bmm_chunk_bwd_IO (A Dout Db : RegionName)
(chunk_size K ngroups SAB SAS SAH SAK SDB SDC SDH SDM SDN SOB SOS SOH SOK
BM BN BCS numCSBlocks : Nat) :
StreamMetaMasked3DKernelIO₂ where
kernel := bbwd_matmul_surface A Dout Db chunk_size (BCS * numCSBlocks) K ngroups
SAB SAS SAH SAK SDB SDC SDH SDM SDN SOB SOS SOH SOK BM BN BCS
inp1 := Dout
inp2 := A
out := Db
nMeta := 0
sty := Fin.elim0
mbuf := Fin.elim0
mwin := Fin.elim0
T := numCSBlocks
B1 := BM * BCS
B2 := BCS * BN
C := BM * BN
outDType := .real
read1 := fun pid₀ pid₁ pid₂ _ t j =>
bbwdDoutAddr pid₁ (pidC pid₂ ngroups) (pidH pid₂ ngroups) (pidM pid₀ K BN)
BM BCS SDB SDC SDH SDN SDM t.val j.val
read2 := fun pid₀ pid₁ pid₂ _ t j =>
bbwdAAddr pid₁ (pidC pid₂ ngroups) (pidH pid₂ ngroups) (pidN pid₀ K BN)
BN BCS chunk_size SAB SAS SAH SAK t.val j.val
write := fun pid₀ pid₁ pid₂ _ j =>
bbwdDbAddr pid₁ (pidC pid₂ ngroups) (pidH pid₂ ngroups) (pidM pid₀ K BN) (pidN pid₀ K BN)
BM BN chunk_size SOB SOS SOH SOK j.val
mask1 := fun _ _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ _ => True
writeMask := fun pid₀ _ _ _ j =>
bbwdWriteOk (pidM pid₀ K BN) (pidN pid₀ K BN) BM BN (BCS * numCSBlocks) K j.val
Referenced model/specification definitions: Python source |
|---|
bmm_chunk_fwdbmm_chunk_fwd_output_summary_generalMathematical execution | Configured model / stage Chunked forward matrix product with reduction extent BK*numKBlocks and the declared group/stride configuration; the public IO skin supplies frame and bounds. | Statement and preconditionsspecification bmm_chunk_fwd_output_summary_general
(A B Out : RegionName) (s : BlockState)
(chunk_size ngroups SAB SAS SAH SAK SBB SBS SBH SBK SOB SOC SOH SOM SON BM BN BK numKBlocks : Nat)
(hBK : 0 < BK)
(hInj : Function.Injective (outOffset (s.pids 1) (pidC (s.pids 2) ngroups) (pidH (s.pids 2) ngroups)
chunk_size SOB SOC SOH SOM SON (pidM (s.pids 0) chunk_size BN) (pidN (s.pids 0) chunk_size BN) BM BN))
(hmlt : ∀ i : Fin BM, rowIndex (pidM (s.pids 0) chunk_size BN) BM i < chunk_size)
(hnlt : ∀ j : Fin BN, colIndex (pidN (s.pids 0) chunk_size BN) BN j < chunk_size)
(hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (bmm_matmul_surface A B Out chunk_size (BK * numKBlocks) ngroups
SAB SAS SAH SAK SBB SBS SBH SBK SOB SOC SOH SOM SON BM BN BK).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := bmm_matmul_surface A B Out chunk_size (BK * numKBlocks) ngroups
SAB SAS SAH SAK SBB SBS SBH SBK SOB SOC SOH SOM SON BM BN BK)
(initialState := s)
(write := fun idx : TileIndex [BM, BN] =>
some (Out, outOffset (s.pids 1) (pidC (s.pids 2) ngroups) (pidH (s.pids 2) ngroups)
chunk_size SOB SOC SOH SOM SON (pidM (s.pids 0) chunk_size BN) (pidN (s.pids 0) chunk_size BN) BM BN idx))
(expected := fun idx : TileIndex [BM, BN] =>
outputCell s A B (s.pids 1) (pidC (s.pids 2) ngroups) (pidH (s.pids 2) ngroups)
(pidM (s.pids 0) chunk_size BN) (pidN (s.pids 0) chunk_size BN)
chunk_size BM BN SAB SAS SAH SAK SBB SBS SBH SBK BK numKBlocks idx)
Referenced model/specification definitions: Python source |
|---|
bmm_chunk_fwdbmm_chunk_fwd_io_correctnessAbstract cast/store rounding | Configured model / stage Chunked forward matrix product with reduction extent BK*numKBlocks and the declared group/stride configuration; the public IO skin supplies frame and bounds. | Statement and preconditionsspecification bmm_chunk_fwd_io_correctness (R : RoundingModel) (A B Out : RegionName)
(chunk_size ngroups SAB SAS SAH SAK SBB SBS SBH SBK SOB SOC SOH SOM SON
BM BN BK numKBlocks : Nat)
(hBK : 0 < BK)
(hmlt : ∀ (pid₀ : Nat) (i : Fin BM), pidM pid₀ chunk_size BN * BM + i.val < chunk_size)
(hnlt : ∀ (pid₀ : Nat) (j : Fin BN), pidN pid₀ chunk_size BN * BN + j.val < chunk_size)
(hOutInj : ∀ pid₀ pid₁ pid₂ : Nat,
Function.Injective (fun idx : TileIndex [BM, BN] =>
bmmOutAddr pid₀ pid₁ pid₂ ngroups chunk_size BM BN SOB SOC SOH SOM SON
idx.1.val idx.2.1.val)) :
bmm_chunk_fwd_IO A B Out chunk_size ngroups SAB SAS SAH SAK SBB SBS SBH SBK
SOB SOC SOH SOM SON BM BN BK numKBlocks ⊨[R]
fun _ _ _ _ xs ys l => bmmStreamSum BM BN BK numKBlocks xs ys l
IO definition: bmm_chunk_fwd_IOdef bmm_chunk_fwd_IO (A B Out : RegionName)
(chunk_size ngroups SAB SAS SAH SAK SBB SBS SBH SBK SOB SOC SOH SOM SON
BM BN BK numKBlocks : Nat) :
StreamMetaMasked3DKernelIO₂ where
kernel := bmm_matmul_surface A B Out chunk_size (BK * numKBlocks) ngroups
SAB SAS SAH SAK SBB SBS SBH SBK SOB SOC SOH SOM SON BM BN BK
inp1 := A
inp2 := B
out := Out
nMeta := 0
sty := Fin.elim0
mbuf := Fin.elim0
mwin := Fin.elim0
T := numKBlocks
B1 := BM * BK
B2 := BK * BN
C := BM * BN
outDType := .real
read1 := fun pid₀ pid₁ pid₂ _ t j =>
bmmAAddr pid₀ pid₁ pid₂ ngroups chunk_size BM BN BK SAB SAS SAH SAK
t.val (j.val / BK) (j.val % BK)
read2 := fun pid₀ pid₁ pid₂ _ t j =>
bmmBAddr pid₀ pid₁ pid₂ ngroups chunk_size BN BK SBB SBS SBH SBK
t.val (j.val / BN) (j.val % BN)
write := fun pid₀ pid₁ pid₂ _ j =>
bmmOutAddr pid₀ pid₁ pid₂ ngroups chunk_size BM BN SOB SOC SOH SOM SON
(j.val / BN) (j.val % BN)
mask1 := fun _ _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ _ => True
writeMask := fun _ _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
bmm_optimizedbmm_o_exec_genuineMathematical execution | Configured model / stage One optimized BMM program under the displayed divisibility, stride, and group-order configuration. The exec-existential proves output values and termination but has no frame conjunct. | Statement and preconditionsspecification bmm_o_exec_genuine
(s : BlockState) (A B O : RegionName)
(M N K TILE_M TILE_N TILE_K GROUP_M : Nat)
(hTN : TILE_N ≤ N) (hTK : 0 < TILE_K)
(hundef : ∀ rg off, s.undef rg off = 0) :
∃ sF, exec (bmm_surface A B O M N K TILE_M TILE_N TILE_K
GROUP_M).toAlgKernel s = some sF
∧ (∀ idx : TileIndex [TILE_M, TILE_N],
bmmOActive s M N TILE_M TILE_N GROUP_M idx →
sF.readMem O (bmmOOffset s M N TILE_M TILE_N GROUP_M idx)
= bmmOOut s A B M N K
(bmmRowG s TILE_M GROUP_M idx.1.val)
(bmmColG s TILE_N GROUP_M idx.2.1.val))
Referenced model/specification definitions: Python source |
|---|
cache_transform_tritondecoding_cache_correctnessMathematical execution | Configured model / stage Decoding is a one-sequence-block specialization; prefill has its separate element-count and region-disjointness premises. Both copy cosine/sine windows, not host cache construction. | Statement and preconditionsspecification decoding_cache_correctness
(cos_cache sin_cache : RegionName) (lengths : Region .nat)
(cos_output sin_output : RegionName)
(cache_stride hidden_stride HIDDEN_DIM NUM_SEQS BLOCK_H : Nat)
(hRegion : cos_output ≠ sin_output) :
decodingCacheIO cos_cache sin_cache lengths cos_output sin_output
cache_stride hidden_stride HIDDEN_DIM NUM_SEQS BLOCK_H
⊨ fun _ _ _ xs ys => (xs, ys)
IO definition: decodingCacheIOdef decodingCacheIO
(cos_cache sin_cache : RegionName) (lengths : Region .nat)
(cos_output sin_output : RegionName)
(cache_stride hidden_stride HIDDEN_DIM NUM_SEQS BLOCK_H : Nat) :
GatherMasked2DKernelIO₂ₓ₂ where
kernel := decoding_cache_one_seq_block cos_cache sin_cache lengths cos_output
sin_output cache_stride hidden_stride HIDDEN_DIM NUM_SEQS BLOCK_H
in1 := cos_cache
in2 := sin_cache
idxbuf := lengths
out1 := cos_output
out2 := sin_output
B := BLOCK_H
N := 1
readx := fun pid₀ _ _ => pid₀
read := fun _ pid₁ ids i =>
ids 0 * cache_stride + (pid₁ * BLOCK_H + i.val) * hidden_stride
write1 := fun pid₀ pid₁ _ i =>
pid₀ * cache_stride + (pid₁ * BLOCK_H + i.val) * hidden_stride
write2 := fun pid₀ pid₁ _ i =>
pid₀ * cache_stride + (pid₁ * BLOCK_H + i.val) * hidden_stride
mask := fun pid₀ _ _ => pid₀ < NUM_SEQS
readMask := fun pid₀ pid₁ _ i =>
pid₀ < NUM_SEQS ∧ pid₁ * BLOCK_H + i.val < HIDDEN_DIM
Referenced model/specification definitions: Python source |
|---|
cache_transform_tritonprefill_cache_correctnessMathematical execution | Configured model / stage Decoding is a one-sequence-block specialization; prefill has its separate element-count and region-disjointness premises. Both copy cosine/sine windows, not host cache construction. | Statement and preconditionsspecification prefill_cache_correctness
(cos_cache sin_cache : RegionName) (cumsum_lengths : Region .nat)
(cos_output sin_output : RegionName)
(cache_stride hidden_stride total_length HIDDEN_DIM N_ELEMENTS BLOCK_SIZE : Nat)
(hRegion : cos_output ≠ sin_output) (hN : 0 < N_ELEMENTS) :
prefillCacheIO cos_cache sin_cache cumsum_lengths cos_output sin_output
cache_stride hidden_stride total_length HIDDEN_DIM N_ELEMENTS BLOCK_SIZE
⊨ fun _ _ _ xs ys => (xs, ys)
IO definition: prefillCacheIOnoncomputable def prefillCacheIO
(cos_cache sin_cache : RegionName) (cumsum_lengths : Region .nat)
(cos_output sin_output : RegionName)
(cache_stride hidden_stride total_length HIDDEN_DIM N_ELEMENTS BLOCK_SIZE : Nat) :
GatherMasked2DKernelIO₂ₓ₂ where
kernel := prefill_cache_kernel cos_cache sin_cache cumsum_lengths cos_output
sin_output cache_stride hidden_stride total_length HIDDEN_DIM N_ELEMENTS BLOCK_SIZE
in1 := cos_cache
in2 := sin_cache
idxbuf := cumsum_lengths
out1 := cos_output
out2 := sin_output
B := HIDDEN_DIM
N := N_ELEMENTS
readx := fun _ _ j => j.val
read := fun pid₀ pid₁ ids i =>
prefillOriSeqIdxOfIds (pid₀ * BLOCK_SIZE + pid₁) N_ELEMENTS ids * cache_stride
+ i.val * hidden_stride
write1 := fun pid₀ pid₁ _ i =>
(pid₀ * BLOCK_SIZE + pid₁) * cache_stride + i.val * hidden_stride
write2 := fun pid₀ pid₁ _ i =>
(pid₀ * BLOCK_SIZE + pid₁) * cache_stride + i.val * hidden_stride
mask := fun _ _ _ => True
readMask := fun pid₀ pid₁ _ _ => pid₀ * BLOCK_SIZE + pid₁ < total_length
Referenced model/specification definitions: Python source |
|---|
chunk_bwd_dqkgchunk_bwd_dqkg_exec_genuineMathematical execution | Configured model / stage Selected backward dq/dk/dg producer with shape/address hypotheses. The exec-existential has output readbacks and termination but no memory-frame conjunct. | Statement and preconditionsspecification chunk_bwd_dqkg_exec_genuine
(q k v h g do_ dh dq dk dg : RegionName)
(s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t : Nat) (scale : ℝ)
(T K V BT BK BV NT : Nat) (s : BlockState)
(hV : 0 < V) (hVB : V ≤ BV)
(hDqDk : dq ≠ dk) (hDqDg : dq ≠ dg) (hDkDg : dk ≠ dg)
(hInj : Function.Injective
(fun i : TileIndex [BT, BK] => cbdQkAddr s s_k_h s_k_t BT BK i)) :
∃ sF, exec (chunk_bwd_dqkg_surface q k v h g do_ dh dq dk dg
s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t scale T K V BT BK BV NT).toAlgKernel s
= some sF
∧ (∀ i : TileIndex [BT, BK],
(s.pids 1 * BT + i.1.val < T ∧ s.pids 0 * BK + i.2.1.val < K) →
sF.readMem dq (cbdQkAddr s s_k_h s_k_t BT BK i)
= dqSpec s g k v h do_ s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t scale
T K V BT BK BV NT i.1.val i.2.1.val)
∧ (∀ i : TileIndex [BT, BK],
(s.pids 1 * BT + i.1.val < T ∧ s.pids 0 * BK + i.2.1.val < K) →
sF.readMem dk (cbdQkAddr s s_k_h s_k_t BT BK i)
= dkSpec s g q v do_ dh s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t scale
T K V BT BK BV NT i.1.val i.2.1.val)
∧ (∀ i : TileIndex [BT], s.pids 1 * BT + i.1.val < T →
sF.readMem dg (cbdDgAddr s T BT i)
= dgSpec s g q k v h do_ dh s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t scale
T K V BT BK BV NT i.1.val)
Referenced model/specification definitions: Python source |
|---|
chunk_cumsum_kernelchunk_cumsum_scalar_output_summary_generalMathematical execution | Configured model / stage Scalar cumulative sum for the modeled per-head sequence and positive chunk size. Read the exact versus stream-rounding contracts separately; launch composition is external. | Statement and preconditionsspecification chunk_cumsum_scalar_output_summary_general
(S O : RegionName) (T BT : Nat) (s : BlockState)
(hSO : O ≠ S) (hBT : 0 < BT) :
-- (1) the full surface lowers to the algorithm layer
(∃ alg, (chunk_cumsum_scalar_surface S O T BT).toAlgorithm? = Except.ok alg) ∧
-- (2) the surface runs to completion (existence / termination)
(∃ sfinal,
exec (chunk_cumsum_scalar_surface S O T BT).toAlgKernel s = some sfinal) ∧
-- (3) standard Realizes_without_Rounding: every in-range O lane holds the genuine global
-- prefix sum `Σ_{m ≤ flat, m < T} S[i_bh·T + m]`, read purely over input
ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_cumsum_scalar_surface S O T BT)
(initialState := s)
(write := fun i : Fin T => some (O, s.pids 0 * T + i.val))
(expected := fun i : Fin T =>
∑ m ∈ (Finset.range T).filter (fun m => m ≤ i.val),
s.readMem S (s.pids 0 * T + m))
Referenced model/specification definitions: Python source |
|---|
chunk_cumsum_kernelchunk_cumsum_scalar_io_correctnessAbstract cast/store rounding | Configured model / stage Scalar cumulative sum for the modeled per-head sequence and positive chunk size. Read the exact versus stream-rounding contracts separately; launch composition is external. | Statement and preconditionsspecification chunk_cumsum_scalar_io_correctness (R : RoundingModel)
(S O : RegionName) (NT BT : Nat) (hSO : O ≠ S) (hBT : 0 < BT) :
chunkCumsumKernelIO S O NT BT ⊨[R]
fun _ _ xs t j => ccStreamSpec NT BT xs t j
IO definition: chunkCumsumKernelIOdef chunkCumsumKernelIO (S O : RegionName) (NT BT : Nat) :
StreamEmitMasked2DKernelIO₁ where
kernel := chunk_cumsum_scalar_surface S O NT BT
inp1 := S
out := O
T := ccNumChunks NT BT
B1 := BT
C := BT
read1 := fun p₀ _ t j => p₀ * NT + (t.val * BT + j.val)
write := fun p₀ _ t j => p₀ * NT + (t.val * BT + j.val)
mask1 := fun _ _ t j => t.val * BT + j.val < NT
writeMask := fun _ _ t j => t.val * BT + j.val < NT
Referenced model/specification definitions: Python source |
|---|
chunk_cumsum_vectorchunk_cumsum_vector_output_summary_generalMathematical execution | Configured model / stage Vector cumulative sum for a single block with an explicit no-address-collision premise. | Statement and preconditionsspecification chunk_cumsum_vector_output_summary_general
(SReg Z : RegionName) (s_s_h s_s_t s_s_d T S BT BS : Nat) (s : BlockState)
(hNoCollision : ∀ idx : TileIndex [BT, BS], singleBlockActive s T S BS idx →
∀ k : TileIndex [BT, BS], singleBlockActive s T S BS k →
singleBlockTileOffset s s_s_h s_s_t s_s_d BS k =
singleBlockTileOffset s s_s_h s_s_t s_s_d BS idx → k = idx) :
ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_cumsum_vector_single_block_surface SReg Z s_s_h s_s_t
s_s_d T S BT BS)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BT, BS] => singleBlockActive s T S BS idx)
(fun idx : TileIndex [BT, BS] =>
(Z, singleBlockTileOffset s s_s_h s_s_t s_s_d BS idx)))
(expected := fun idx : TileIndex [BT, BS] =>
singleBlockCumsumVectorClosed s SReg s_s_h s_s_t s_s_d T S BS idx)
Referenced model/specification definitions: Python source |
|---|
chunk_cumsum_vectorchunk_cumsum_vector_block_store_io_correctnessMathematical execution | Precomputed-input slice Copies precomputed block-cumsum tile BC into Z; it does not compute the cumulative sum. | Statement and preconditionsspecification chunk_cumsum_vector_block_store_io_correctness (BC Z : RegionName)
(s_s_h s_s_t s_s_d T S BT BS : Nat)
(hOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BT, BS] =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t
+ (p₀ * BS + idx.2.1.val) * s_s_d)) :
blockStoreIO BC Z s_s_h s_s_t s_s_d T S BT BS
⊨ fun _p₀ _p₁ xs idx => xs idx
IO definition: blockStoreIOdef blockStoreIO (BC Z : RegionName) (s_s_h s_s_t s_s_d T S BT BS : Nat) :
Masked3DTileKernelIO₁ where
kernel := chunk_cumsum_vector_store_slice BC Z s_s_h s_s_t s_s_d T S BT BS
inp := BC
out := Z
shape := [BT, BS]
read := fun p₀ p₁ p₂ idx =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t + (p₀ * BS + idx.2.1.val) * s_s_d
write := fun p₀ p₁ p₂ idx =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t + (p₀ * BS + idx.2.1.val) * s_s_d
mask := fun p₀ _p₁ p₂ idx =>
p₂ * BT + idx.1.val < T ∧ p₀ * BS + idx.2.1.val < S
Referenced model/specification definitions: Python source |
|---|
chunk_cumsum_vectorchunk_cumsum_vector_block_store_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract store-rounding version of the BC-to-Z copy slice; cumulative-sum production is outside this contract. | Statement and preconditionsspecification chunk_cumsum_vector_block_store_io_correctnessR (R : RoundingModel)
(BC Z : RegionName) (s_s_h s_s_t s_s_d T S BT BS : Nat)
(hOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BT, BS] =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t
+ (p₀ * BS + idx.2.1.val) * s_s_d)) :
blockStoreIO BC Z s_s_h s_s_t s_s_d T S BT BS
⊨[R, FloatDType.real] fun _p₀ _p₁ xs idx => xs idx
IO definition: blockStoreIOdef blockStoreIO (BC Z : RegionName) (s_s_h s_s_t s_s_d T S BT BS : Nat) :
Masked3DTileKernelIO₁ where
kernel := chunk_cumsum_vector_store_slice BC Z s_s_h s_s_t s_s_d T S BT BS
inp := BC
out := Z
shape := [BT, BS]
read := fun p₀ p₁ p₂ idx =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t + (p₀ * BS + idx.2.1.val) * s_s_d
write := fun p₀ p₁ p₂ idx =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t + (p₀ * BS + idx.2.1.val) * s_s_d
mask := fun p₀ _p₁ p₂ idx =>
p₂ * BT + idx.1.val < T ∧ p₀ * BS + idx.2.1.val < S
Referenced model/specification definitions: Python source |
|---|
chunk_delta_fwdchunk_delta_fwd_exec_genuineMathematical execution | Configured model / stage Forward state/v_new/final-state recurrence for the selected state flags, with explicit output readbacks. The exec-existential does not include a frame. | Statement and preconditionsspecification chunk_delta_fwd_exec_genuine
(k v d v_new h initial_state final_state : RegionName)
(USE_INITIAL_STATE STORE_FINAL_STATE : Bool)
(s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d s_h_h s_h_t
H T K V BT BC BK BV NT : Nat)
(s : BlockState) (hpids0 : s.pids 0 = 0)
(hBC : BC = BT) (hBT : 0 < BT) (hBK : BK ≤ K) (hTNT : NT * BT ≤ T)
(hVod : 0 < s_vo_d)
(hHBlock : (BK - 1) * s_h_t + BV ≤ K * V) (hVBlock : BV * s_vo_d ≤ s_vo_t)
(hVk : v_new ≠ k) (hVv : v_new ≠ v) (hVd : v_new ≠ d) (hHv : h ≠ v_new)
(hHk : h ≠ k) (hHv2 : h ≠ v) (hHd : h ≠ d)
(hFh : final_state ≠ h) (hFv : final_state ≠ v_new) (hFk : final_state ≠ k)
(hFv2 : final_state ≠ v) (hFd : final_state ≠ d)
(hInjV : ∀ i_t : Fin NT, Function.Injective
(fun idx : TileIndex [BC, BV] =>
cdfVNewAddr s s_vo_h s_vo_t s_vo_d BT BC BV i_t.val idx))
(hInjH : ∀ i_t : Fin NT, Function.Injective
(fun idx : TileIndex [BK, BV] => hOffset s i_t.val s_h_h s_h_t K V BK BV idx))
(hInjF : Function.Injective
(fun idx : TileIndex [BK, BV] => finalStateOffset s K V BK BV idx)) :
∃ sF, exec (chunk_delta_rule_fwd_h_surface k v d v_new h initial_state final_state
s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d s_h_h s_h_t H T K V BT BC BK BV NT
USE_INITIAL_STATE STORE_FINAL_STATE).toAlgKernel s = some sF
∧ (∀ j : Fin NT, ∀ idx : TileIndex [BK, BV], active s K V BK BV idx →
sF.readMem h (hOffset s j.val s_h_h s_h_t K V BK BV idx)
= hValue s k v d initial_state s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d
K V BT BV BK USE_INITIAL_STATE j.val idx)
∧ (∀ j : Fin NT, ∀ idx : TileIndex [BC, BV],
vNewActive s j.val 0 T V BT BC BV idx →
sF.readMem v_new (cdfVNewAddr s s_vo_h s_vo_t s_vo_d BT BC BV j.val idx)
= vNewSpec s k v d initial_state s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d
K V BT BV BK BC USE_INITIAL_STATE j.val idx)
∧ (STORE_FINAL_STATE = Bool.true →
∀ idx : TileIndex [BK, BV], active s K V BK BV idx →
sF.readMem final_state (finalStateOffset s K V BK BV idx)
= finalValue s k v d initial_state s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d
K V BT BV BK USE_INITIAL_STATE NT idx)
Referenced model/specification definitions: Python source |
|---|
chunk_delta_fwdchunk_delta_h_state_store_io_correctnessMathematical execution | Precomputed-input slice Writes a precomputed HPre tile to one recurrent-state slot; no recurrence producer is inside this IO kernel. | Statement and preconditionsspecification chunk_delta_h_state_store_io_correctness
(HPre h : RegionName) (i_t s_h_h s_h_t K V BK BV : Nat)
(hInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BK, BV] =>
p₂ * s_h_h + i_t * K * V + (p₀ * BK + idx.1.val) * s_h_t
+ (p₁ * BV + idx.2.1.val))) :
h_stateIO HPre h i_t s_h_h s_h_t K V BK BV
⊨ fun _p₀ _p₁ xs idx => xs idx
IO definition: h_stateIOdef h_stateIO (HPre h : RegionName) (i_t s_h_h s_h_t K V BK BV : Nat) :
Masked3DTileKernelIO₁ where
kernel := chunk_delta_h_state_store_slice HPre h i_t s_h_h s_h_t K V BK BV
inp := HPre
out := h
shape := [BK, BV]
read := fun p₀ p₁ p₂ idx =>
p₂ * s_h_h + i_t * K * V + (p₀ * BK + idx.1.val) * s_h_t
+ (p₁ * BV + idx.2.1.val)
write := fun p₀ p₁ p₂ idx =>
p₂ * s_h_h + i_t * K * V + (p₀ * BK + idx.1.val) * s_h_t
+ (p₁ * BV + idx.2.1.val)
mask := fun p₀ p₁ _p₂ idx =>
p₀ * BK + idx.1.val < K ∧ p₁ * BV + idx.2.1.val < V
Referenced model/specification definitions: Python source |
|---|
chunk_delta_fwdchunk_delta_h_state_store_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding contract for the precomputed HPre state-store slice. | Statement and preconditionsspecification chunk_delta_h_state_store_io_correctnessR
(R : RoundingModel) (HPre h : RegionName) (i_t s_h_h s_h_t K V BK BV : Nat)
(hInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BK, BV] =>
p₂ * s_h_h + i_t * K * V + (p₀ * BK + idx.1.val) * s_h_t
+ (p₁ * BV + idx.2.1.val))) :
h_stateIO HPre h i_t s_h_h s_h_t K V BK BV
⊨[R, FloatDType.real] fun _p₀ _p₁ xs idx => xs idx
IO definition: h_stateIOdef h_stateIO (HPre h : RegionName) (i_t s_h_h s_h_t K V BK BV : Nat) :
Masked3DTileKernelIO₁ where
kernel := chunk_delta_h_state_store_slice HPre h i_t s_h_h s_h_t K V BK BV
inp := HPre
out := h
shape := [BK, BV]
read := fun p₀ p₁ p₂ idx =>
p₂ * s_h_h + i_t * K * V + (p₀ * BK + idx.1.val) * s_h_t
+ (p₁ * BV + idx.2.1.val)
write := fun p₀ p₁ p₂ idx =>
p₂ * s_h_h + i_t * K * V + (p₀ * BK + idx.1.val) * s_h_t
+ (p₁ * BV + idx.2.1.val)
mask := fun p₀ p₁ _p₂ idx =>
p₀ * BK + idx.1.val < K ∧ p₁ * BV + idx.2.1.val < V
Referenced model/specification definitions: Python source |
|---|
chunk_gate_recurrencechunk_gate_recurrence_output_summary_generalMathematical execution | Configured model / stage Bundle of forward/backward recurrence and individual step/store facts. Each conjunct has its own target; it is not a single whole-program IO contract. | Statement and preconditionsspecification chunk_gate_recurrence_output_summary_general
(AccPrev S D O LastKv : RegionName) (HAS_LAST_KV : Bool)
(DaccPrev DaccTail DS DI DG DL : RegionName) (t_rel : Nat)
(NUM_HEAD NUM_BLOCK NUM_K NUM_V D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V : Nat)
(s : BlockState)
-- forced by `BlockState.scatter_readback_nd`: the forward initial tile store
-- must not have two lanes colliding on one `O` cell.
(hOutInj0 : Function.Injective
(fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
outOffset s NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V idx))
-- same, for the forward loop body's store into `O` chunk `t_rel + 1`.
(hOutInjStep : Function.Injective
(fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
forwardStepTileOffset s (t_rel + 1) NUM_BLOCK D_MODEL_K D_MODEL_V
BLOCK_MODEL_K BLOCK_MODEL_V idx))
-- same, for the reverse loop body's store into `DI` chunk `t_rel`.
(hDIInj : Function.Injective
(fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
timeTileOffset s t_rel NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V idx))
-- same, for the post-loop store into the chunk-free `DL` tile.
(hDLInj : Function.Injective
(fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
accOffset s D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V idx))
-- forced by `bwdClosed_step`: the reverse loop body's gate/`DS` index is
-- `t_rel + 1`, and peeling it off the fold needs it to be a real chunk.
-- Python's loop runs `t_rel = NUM_BLOCK-2 … 0`, so `t_rel + 1 ≤ NUM_BLOCK-1`.
-- (`bwdClosed_step` at the post-loop index `0` needs `0 < NUM_BLOCK`; that is
-- implied by `hBwdIdx`, so it is not stated separately.)
(hBwdIdx : t_rel + 1 < NUM_BLOCK)
-- the forward carry: `acc` entering loop iteration `t_rel` is the `t_rel`-step
-- fold. Assumed — the `range(NUM_BLOCK-1)` fold is the trusted boundary.
(hAcc : ∀ idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V],
s.readMem AccPrev
(accOffset s D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V idx)
= fwdClosed s S D LastKv HAS_LAST_KV NUM_BLOCK D_MODEL_K D_MODEL_V
BLOCK_MODEL_K BLOCK_MODEL_V t_rel idx)
-- the reverse carry: `Dacc` entering the iteration that writes chunk `t_rel`
-- is the reverse fold at gate index `t_rel + 2`. Assumed for the same reason;
-- at the loop's first iteration (`t_rel = NUM_BLOCK - 2`) it is discharged by
-- `bwdClosed_top`, matching Python's `Dacc = tl.zeros(...)` seed.
(hDacc : ∀ idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V],
s.readMem DaccPrev
(accOffset s D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V idx)
= bwdClosed s DS D NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V (t_rel + 2) idx)
-- the reverse carry as the loop *exits* (gate index `1`), feeding the
-- post-loop `DL` step. A separate region from `DaccPrev` on purpose: pinning
-- one buffer to two different fold stages would silently constrain `t_rel`.
(hDaccTail : ∀ idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V],
s.readMem DaccTail
(accOffset s D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V idx)
= bwdClosed s DS D NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V 1 idx) :
(∃ alg, (chunk_gate_recurrence_fwd_surface S D O LastKv
NUM_HEAD NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V
Bool.true).toAlgorithm? = Except.ok alg) ∧
(∃ alg, (chunk_gate_recurrence_fwd_surface S D O LastKv
NUM_HEAD NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V
Bool.false).toAlgorithm? = Except.ok alg) ∧
((chunk_gate_recurrence_bwd_surface S D DI DG DL DS
NUM_HEAD NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V).toAlgorithm? =
Except.ok
(chunk_gate_recurrence_bwd_surface S D DI DG DL DS
NUM_HEAD NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V).toAlgKernel) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_gate_recurrence_initial_last_kv_store_slice LastKv O
NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
some (O, outOffset s NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V idx))
(expected := fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
fwdClosed s S D LastKv Bool.true NUM_BLOCK D_MODEL_K D_MODEL_V
BLOCK_MODEL_K BLOCK_MODEL_V 0 idx)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_gate_recurrence_initial_zero_store_slice O NUM_BLOCK
D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
some (O, outOffset s NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V idx))
(expected := fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
fwdClosed s S D LastKv Bool.false NUM_BLOCK D_MODEL_K D_MODEL_V
BLOCK_MODEL_K BLOCK_MODEL_V 0 idx)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_gate_recurrence_forward_step_store_slice AccPrev S D O
t_rel NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
some (O, forwardStepTileOffset s (t_rel + 1) NUM_BLOCK D_MODEL_K
D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V idx))
(expected := fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
fwdClosed s S D LastKv HAS_LAST_KV NUM_BLOCK D_MODEL_K D_MODEL_V
BLOCK_MODEL_K BLOCK_MODEL_V (t_rel + 1) idx)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_gate_recurrence_bwd_dacc_step_DI_store_slice DaccPrev
DS D DI t_rel NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
some (DI, timeTileOffset s t_rel NUM_BLOCK D_MODEL_K D_MODEL_V
BLOCK_MODEL_K BLOCK_MODEL_V idx))
(expected := fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
bwdClosed s DS D NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V (t_rel + 1) idx)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_gate_recurrence_bwd_dg_step_store_slice DaccPrev DS
S D DG t_rel NUM_BLOCK NUM_K NUM_V D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V)
(initialState := s)
(write := fun _ : PUnit =>
some (DG, bwdDGOffset s t_rel NUM_BLOCK NUM_K NUM_V))
(expected := fun _ =>
bwdDGClosed s DS S D t_rel NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_gate_recurrence_bwd_DL_store_slice DaccTail DS D DL
0 NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
some (DL, accOffset s D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V idx))
(expected := fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
bwdClosed s DS D NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K
BLOCK_MODEL_V 0 idx))
Referenced model/specification definitions: Python source |
|---|
chunk_gate_recurrencechunk_gate_recurrence_forward_store_io_correctnessMathematical execution | Precomputed-input slice Copies a supplied accumulator tile Acc to O; the recurrence computing Acc is outside this IO contract. | Statement and preconditionsspecification chunk_gate_recurrence_forward_store_io_correctness
(Acc O : RegionName) (NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V : Nat)
(hOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
p₀ * NUM_BLOCK * D_MODEL_K * D_MODEL_V + p₁ * D_MODEL_V * BLOCK_MODEL_K
+ idx.1.val * D_MODEL_V + p₂ * BLOCK_MODEL_V + idx.2.1.val)) :
fwdStoreIO Acc O NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V ⊨ fun _p₀ _p₁ xs idx => xs idx
IO definition: fwdStoreIOdef fwdStoreIO (Acc O : RegionName) (NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V : Nat) : Masked3DTileKernelIO₁ where
kernel := chunk_gate_recurrence_forward_store_slice Acc O NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V
inp := Acc
out := O
shape := [BLOCK_MODEL_K, BLOCK_MODEL_V]
read := fun p₀ p₁ p₂ idx =>
p₀ * D_MODEL_K * D_MODEL_V + p₁ * D_MODEL_V * BLOCK_MODEL_K
+ idx.1.val * D_MODEL_V + p₂ * BLOCK_MODEL_V + idx.2.1.val
write := fun p₀ p₁ p₂ idx =>
p₀ * NUM_BLOCK * D_MODEL_K * D_MODEL_V + p₁ * D_MODEL_V * BLOCK_MODEL_K
+ idx.1.val * D_MODEL_V + p₂ * BLOCK_MODEL_V + idx.2.1.val
mask := fun _p₀ _p₁ _p₂ _ => True
Referenced model/specification definitions: Python source |
|---|
chunk_gate_recurrencechunk_gate_recurrence_forward_store_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding contract for the supplied-accumulator forward-store slice. | Statement and preconditionsspecification chunk_gate_recurrence_forward_store_io_correctnessR
(R : RoundingModel) (Acc O : RegionName)
(NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V : Nat)
(hOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_MODEL_K, BLOCK_MODEL_V] =>
p₀ * NUM_BLOCK * D_MODEL_K * D_MODEL_V + p₁ * D_MODEL_V * BLOCK_MODEL_K
+ idx.1.val * D_MODEL_V + p₂ * BLOCK_MODEL_V + idx.2.1.val)) :
fwdStoreIO Acc O NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V
⊨[R, FloatDType.real] fun _p₀ _p₁ xs idx => xs idx
IO definition: fwdStoreIOdef fwdStoreIO (Acc O : RegionName) (NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V : Nat) : Masked3DTileKernelIO₁ where
kernel := chunk_gate_recurrence_forward_store_slice Acc O NUM_BLOCK D_MODEL_K D_MODEL_V BLOCK_MODEL_K BLOCK_MODEL_V
inp := Acc
out := O
shape := [BLOCK_MODEL_K, BLOCK_MODEL_V]
read := fun p₀ p₁ p₂ idx =>
p₀ * D_MODEL_K * D_MODEL_V + p₁ * D_MODEL_V * BLOCK_MODEL_K
+ idx.1.val * D_MODEL_V + p₂ * BLOCK_MODEL_V + idx.2.1.val
write := fun p₀ p₁ p₂ idx =>
p₀ * NUM_BLOCK * D_MODEL_K * D_MODEL_V + p₁ * D_MODEL_V * BLOCK_MODEL_K
+ idx.1.val * D_MODEL_V + p₂ * BLOCK_MODEL_V + idx.2.1.val
mask := fun _p₀ _p₁ _p₂ _ => True
Referenced model/specification definitions: Python source |
|---|
chunk_gated_attentionchunk_gated_attention_cum_slice_output_summary_generalMathematical execution | Configured model / stage Cumulative-sum and gated-state step slices plus recurrence identities. Previous-state and step premises remain explicit; no composed end-to-end attention contract. | Statement and preconditionsspecification chunk_gated_attention_cum_slice_output_summary_general
(SReg GCum K V G H0 HPrev HOut : RegionName) (USE_INITIAL_STATE : Bool)
(s_s_h s_s_t s_s_d T S BT BS
m s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d KSize VSize BK BV : Nat)
(s : BlockState)
(hCumInj : Function.Injective
(fun idx : TileIndex [BT, BS] => tileOffset s s_s_h s_s_t s_s_d BT BS idx))
(hStateInj : Function.Injective
(fun idx : TileIndex [BK, BV] => finalStateOffset s KSize VSize BK BV idx))
(hPrevV : ∀ idx : TileIndex [BK, BV],
s.readMem HPrev (finalStateOffset s KSize VSize BK BV idx)
= hClosed s K V G H0 Bool.false USE_INITIAL_STATE
s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d KSize VSize BT BK BV m idx)
(hPrevK : ∀ idx : TileIndex [BK, BV],
s.readMem HPrev (finalStateOffset s KSize VSize BK BV idx)
= hClosed s K V G H0 Bool.true USE_INITIAL_STATE
s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d KSize VSize BT BK BV m idx) :
-- (1) the `fwd_pre` cumsum body realizes the causal intra-chunk cumsum
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_gated_attention_cum_compute_slice SReg GCum s_s_h s_s_t
s_s_d T S BT BS)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BT, BS] => cumSurfaceActive s T S BT BS idx)
(fun idx : TileIndex [BT, BS] =>
(GCum, tileOffset s s_s_h s_s_t s_s_d BT BS idx)))
(expected := fun idx : TileIndex [BT, BS] =>
cumComputeStoreValue s SReg s_s_h s_s_t s_s_d T S BT BS idx)) ∧
-- (2) one `GATEK = false` recurrence body realizes `hClosed (m+1)`
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_gated_attention_h_step_gatev_slice K V G HPrev HOut
m s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d KSize VSize BT BK BV)
(initialState := s)
(write := fun idx : TileIndex [BK, BV] =>
some (HOut, finalStateOffset s KSize VSize BK BV idx))
(expected := fun idx : TileIndex [BK, BV] =>
hClosed s K V G H0 Bool.false USE_INITIAL_STATE
s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d KSize VSize BT BK BV (m + 1) idx)) ∧
-- (3) one `GATEK = true` recurrence body realizes `hClosed (m+1)`
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_gated_attention_h_step_gatek_slice K V G HPrev HOut
m s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d KSize VSize BT BK BV)
(initialState := s)
(write := fun idx : TileIndex [BK, BV] =>
some (HOut, finalStateOffset s KSize VSize BK BV idx))
(expected := fun idx : TileIndex [BK, BV] =>
hClosed s K V G H0 Bool.true USE_INITIAL_STATE
s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d KSize VSize BT BK BV (m + 1) idx)) ∧
-- (4) the recurrence's base case
(∀ (GATEK : Bool) (idx : TileIndex [BK, BV]),
hClosed s K V G H0 GATEK USE_INITIAL_STATE
s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d KSize VSize BT BK BV 0 idx
= hSeed s H0 USE_INITIAL_STATE KSize VSize BK BV idx)
Referenced model/specification definitions: Python source |
|---|
chunk_gated_attentionchunk_gated_attention_state_stores_io_correctnessMathematical execution | Precomputed-input slice Copies supplied BH and BHFinal tiles to intermediate/final state; state production is outside these IO kernels. | Statement and preconditionsspecification chunk_gated_attention_state_stores_io_correctness
(BH H BHFinal Ht : RegionName)
(i_t s_h_h s_h_t s_h_d KSize VSize BK BV : Nat)
(hInj1 : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BK, BV] =>
p₂ * s_h_h + i_t * KSize * VSize + (p₁ * BK + idx.1.val) * s_h_t
+ (p₀ * BV + idx.2.1.val) * s_h_d))
(hInj2 : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BK, BV] =>
p₂ * KSize * VSize + (p₁ * BK + idx.1.val) * VSize
+ (p₀ * BV + idx.2.1.val))) :
(h_stateIO BH H i_t s_h_h s_h_t s_h_d KSize VSize BK BV
⊨ fun _p₀ _p₁ xs idx => xs idx) ∧
(final_stateIO BHFinal Ht KSize VSize BK BV
⊨ fun _p₀ _p₁ xs idx => xs idx)
IO definition: h_stateIOdef h_stateIO (BH H : RegionName) (i_t s_h_h s_h_t s_h_d KSize VSize BK BV : Nat) :
Masked3DTileKernelIO₁ where
kernel := chunk_gated_attention_h_state_store_slice BH H i_t s_h_h s_h_t s_h_d KSize VSize BK BV
inp := BH
out := H
shape := [BK, BV]
read := fun p₀ p₁ p₂ idx => p₂ * s_h_h + i_t * KSize * VSize + (p₁ * BK + idx.1.val) * s_h_t
+ (p₀ * BV + idx.2.1.val) * s_h_d
write := fun p₀ p₁ p₂ idx => p₂ * s_h_h + i_t * KSize * VSize + (p₁ * BK + idx.1.val) * s_h_t
+ (p₀ * BV + idx.2.1.val) * s_h_d
mask := fun p₀ p₁ _p₂ idx =>
p₁ * BK + idx.1.val < KSize ∧ p₀ * BV + idx.2.1.val < VSize
IO definition: final_stateIOdef final_stateIO (BHFinal Ht : RegionName) (KSize VSize BK BV : Nat) :
Masked3DTileKernelIO₁ where
kernel := chunk_gated_attention_final_state_store_slice BHFinal Ht KSize VSize BK BV
inp := BHFinal
out := Ht
shape := [BK, BV]
read := fun p₀ p₁ p₂ idx => p₂ * KSize * VSize + (p₁ * BK + idx.1.val) * VSize
+ (p₀ * BV + idx.2.1.val)
write := fun p₀ p₁ p₂ idx => p₂ * KSize * VSize + (p₁ * BK + idx.1.val) * VSize
+ (p₀ * BV + idx.2.1.val)
mask := fun p₀ p₁ _p₂ idx =>
p₁ * BK + idx.1.val < KSize ∧ p₀ * BV + idx.2.1.val < VSize
Referenced model/specification definitions: Python source |
|---|
chunk_gated_attentionchunk_gated_attention_state_stores_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding contracts for the supplied BH/BHFinal state stores. | Statement and preconditionsspecification chunk_gated_attention_state_stores_io_correctnessR
(R : RoundingModel) (BH H BHFinal Ht : RegionName)
(i_t s_h_h s_h_t s_h_d KSize VSize BK BV : Nat)
(hInj1 : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BK, BV] =>
p₂ * s_h_h + i_t * KSize * VSize + (p₁ * BK + idx.1.val) * s_h_t
+ (p₀ * BV + idx.2.1.val) * s_h_d))
(hInj2 : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BK, BV] =>
p₂ * KSize * VSize + (p₁ * BK + idx.1.val) * VSize
+ (p₀ * BV + idx.2.1.val))) :
(h_stateIO BH H i_t s_h_h s_h_t s_h_d KSize VSize BK BV
⊨[R, FloatDType.real] fun _p₀ _p₁ xs idx => xs idx) ∧
(final_stateIO BHFinal Ht KSize VSize BK BV
⊨[R, FloatDType.real] fun _p₀ _p₁ xs idx => xs idx)
IO definition: h_stateIOdef h_stateIO (BH H : RegionName) (i_t s_h_h s_h_t s_h_d KSize VSize BK BV : Nat) :
Masked3DTileKernelIO₁ where
kernel := chunk_gated_attention_h_state_store_slice BH H i_t s_h_h s_h_t s_h_d KSize VSize BK BV
inp := BH
out := H
shape := [BK, BV]
read := fun p₀ p₁ p₂ idx => p₂ * s_h_h + i_t * KSize * VSize + (p₁ * BK + idx.1.val) * s_h_t
+ (p₀ * BV + idx.2.1.val) * s_h_d
write := fun p₀ p₁ p₂ idx => p₂ * s_h_h + i_t * KSize * VSize + (p₁ * BK + idx.1.val) * s_h_t
+ (p₀ * BV + idx.2.1.val) * s_h_d
mask := fun p₀ p₁ _p₂ idx =>
p₁ * BK + idx.1.val < KSize ∧ p₀ * BV + idx.2.1.val < VSize
IO definition: final_stateIOdef final_stateIO (BHFinal Ht : RegionName) (KSize VSize BK BV : Nat) :
Masked3DTileKernelIO₁ where
kernel := chunk_gated_attention_final_state_store_slice BHFinal Ht KSize VSize BK BV
inp := BHFinal
out := Ht
shape := [BK, BV]
read := fun p₀ p₁ p₂ idx => p₂ * KSize * VSize + (p₁ * BK + idx.1.val) * VSize
+ (p₀ * BV + idx.2.1.val)
write := fun p₀ p₁ p₂ idx => p₂ * KSize * VSize + (p₁ * BK + idx.1.val) * VSize
+ (p₀ * BV + idx.2.1.val)
mask := fun p₀ p₁ _p₂ idx =>
p₁ * BK + idx.1.val < KSize ∧ p₀ * BV + idx.2.1.val < VSize
Referenced model/specification definitions: Python source |
|---|
chunk_gla_fwdchunk_gla_fwd_o_exec_genuineMathematical execution | Configured model / stage The output stage consumes A and h from its input memory; Python's separate A producers are not proved by this headline. The exec-existential omits a frame. | Statement and preconditionsspecification chunk_gla_fwd_o_exec_genuine
(q v g h o A : RegionName)
(s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t : Nat)
(scale : ℝ) (T K V BT BK BV : Nat) (s : BlockState)
(hInj : Function.Injective
(fun idx : TileIndex [BT, BV] => outOffset s s_v_h s_v_t BT BV idx)) :
∃ sF, exec (chunk_gla_fwd_o_surface q v g h o A s_k_h s_k_t s_v_h s_v_t
s_h_h s_h_t scale T K V BT BK BV).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BT, BV],
(tIndex s BT idx.1.val < T ∧ vIndex s BV idx.2.1.val < V) →
sF.readMem o (outOffset s s_v_h s_v_t BT BV idx)
= cgfOutput s q v g h A s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t scale
T K V BT BK BV idx.1.val idx.2.1.val
Referenced model/specification definitions: Python source |
|---|
chunk_gla_simplechunk_gla_simple_output_summary_generalMathematical execution | Configured model / stage Selected forward output stage under the displayed tile-size and injective-address premises. Input h is supplied by the surrounding algorithm; no whole-wrapper state construction claim. | Statement and preconditionsspecification chunk_gla_simple_output_summary_general
(q k v h g o : RegionName)
(s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t : Nat)
(scale : ℝ) (T K V BT BK BV : Nat) (s : BlockState)
(hKBK : K = BK) (hBK : 0 < BK) (hBT : 0 < BT)
(hundef : ∀ rg off, s.undef rg off = 0)
(hInj : Function.Injective
(fun idx : TileIndex [BT, BV] => outOffset s s_v_h s_v_t BT BV idx)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := chunk_gla_simple_fwd_surface q k v h g o s_k_h s_k_t s_v_h s_v_t
s_h_h s_h_t scale T K V BT BK BV)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BT, BV] => active s T V BT BV idx)
(fun idx => (o, outOffset s s_v_h s_v_t BT BV idx)))
(expected := fun idx : TileIndex [BT, BV] =>
glaOutput s q k v h g s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t
scale T K V BT BK BV idx.1 idx.2.1)
Referenced model/specification definitions: Python source |
|---|
chunk_gla_simplechunk_gla_simple_io_correctnessAbstract cast/store rounding | Configured model / stage Selected forward output stage under the displayed tile-size and injective-address premises. Input h is supplied by the surrounding algorithm; no whole-wrapper state construction claim. | Statement and preconditionsspecification chunk_gla_simple_io_correctness (R : RoundingModel)
(q k v h g o : RegionName)
(s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t : Nat) (scale : ℝ)
(Tt K V BT BK BV : Nat)
(hKBK : K = BK) (hBK : 0 < BK) (hBT : 0 < BT)
(hInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BT, BV] =>
p₂ * s_v_h + (p₁ * BT + idx.1.val) * s_v_t + (p₀ * BV + idx.2.1.val) * 1)) :
chunkGlaSimpleIO q k v h g o s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t scale
Tt K V BT BK BV ⊨[R]
fun _ _ _ xs ys zs ws vs j =>
chunkGlaSimpleIOOutSpec scale BT BK BV xs ys zs ws vs j
IO definition: chunkGlaSimpleIOdef chunkGlaSimpleIO (q k v h g o : RegionName)
(s_k_h s_k_t s_v_h s_v_t s_h_h s_h_t : Nat)
(scale : ℝ) (Tt K V BT BK BV : Nat) : StreamMasked3DKernelIO₅ where
kernel := chunk_gla_simple_fwd_surface q k v h g o s_k_h s_k_t s_v_h s_v_t
s_h_h s_h_t scale Tt K V BT BK BV
inp1 := q
inp2 := k
inp3 := h
inp4 := g
inp5 := v
out := o
T := 1
B1 := BT * BK
B2 := BK * BT
B3 := BK * BV
B4 := BT
B5 := BT * BV
C := BT * BV
outDType := .real
read1 := fun _ p₁ p₂ _ j => p₂ * s_k_h + (p₁ * BT + j.val / BK) * s_k_t + j.val % BK
read2 := fun _ p₁ p₂ _ j => p₂ * s_k_h + (j.val / BT) * 1 + (p₁ * BT + j.val % BT) * s_k_t
read3 := fun p₀ p₁ p₂ _ j =>
p₂ * s_h_h + p₁ * K * V + (j.val / BV) * s_h_t + (p₀ * BV + j.val % BV) * 1
read4 := fun _ p₁ p₂ _ j => p₂ * Tt + (p₁ * BT + j.val) * 1
read5 := fun p₀ p₁ p₂ _ j =>
p₂ * s_v_h + (p₁ * BT + j.val / BV) * s_v_t + (p₀ * BV + j.val % BV) * 1
write := fun p₀ p₁ p₂ j =>
p₂ * s_v_h + (p₁ * BT + j.val / BV) * s_v_t + (p₀ * BV + j.val % BV) * 1
mask1 := fun _ p₁ _ _ j => p₁ * BT + j.val / BK < Tt ∧ j.val % BK < K
mask2 := fun _ p₁ _ _ j => j.val / BT < K ∧ p₁ * BT + j.val % BT < Tt
mask3 := fun p₀ _ _ _ j => j.val / BV < K ∧ p₀ * BV + j.val % BV < V
mask4 := fun _ p₁ _ _ j => p₁ * BT + j.val < Tt
mask5 := fun p₀ p₁ _ _ j => p₁ * BT + j.val / BV < Tt ∧ p₀ * BV + j.val % BV < V
writeMask := fun p₀ p₁ _ j => p₁ * BT + j.val / BV < Tt ∧ p₀ * BV + j.val % BV < V
Referenced model/specification definitions: Python source |
|---|
chunk_linear_attncla_fwd_h_exec_genuineMathematical execution | Configured model / stage Forward h/ht and backward dh stage contracts. Separate Python output/dqkv stages and host orchestration are not covered; both exec-existentials omit a frame. | Statement and preconditionsspecification cla_fwd_h_exec_genuine
(k v h h0 ht : RegionName)
(s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d s_h_h s_h_t : Nat)
(T K V BT BK BV NT : Nat) (USE_INITIAL_STATE STORE_FINAL_STATE : Bool)
(s : BlockState)
(hHk : h ≠ k) (hHv : h ≠ v) (hHtH : ht ≠ h)
(hσ : BV ≤ s_h_t) (hFit : (K - 1) * s_h_t + V ≤ K * V) (hBVV : BV ≤ V) :
∃ sF, exec (cla_fwd_h_surface k v h h0 ht s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t
s_vo_d s_h_h s_h_t T K V BT BK BV NT
USE_INITIAL_STATE STORE_FINAL_STATE).toAlgKernel s = some sF
∧ (∀ t (idx : TileIndex [BK, BV]), t < NT → claActive s K V BK BV idx →
sF.readMem h (claHOffset s s_h_h s_h_t K V BK BV t idx)
= claHState s k v h0 USE_INITIAL_STATE s_qk_h s_qk_t s_qk_d s_vo_h
s_vo_t s_vo_d T K V BT BK BV t idx.1.val idx.2.1.val)
∧ (STORE_FINAL_STATE = Bool.true →
∀ idx : TileIndex [BK, BV], claActive s K V BK BV idx →
sF.readMem ht (claHtOffset s K V BK BV idx)
= claHState s k v h0 USE_INITIAL_STATE s_qk_h s_qk_t s_qk_d s_vo_h
s_vo_t s_vo_d T K V BT BK BV NT idx.1.val idx.2.1.val)
Referenced model/specification definitions: Python source |
|---|
chunk_linear_attncla_bwd_dh_exec_genuineMathematical execution | Configured model / stage Forward h/ht and backward dh stage contracts. Separate Python output/dqkv stages and host orchestration are not covered; both exec-existentials omit a frame. | Statement and preconditionsspecification cla_bwd_dh_exec_genuine
(q do_ dh : RegionName)
(s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d s_h_h s_h_t : Nat)
(scale : ℝ) (T K V BT BK BV NT : Nat) (s : BlockState)
(hDq : dh ≠ q) (hDdo : dh ≠ do_)
(hσ : BV ≤ s_h_t) (hFit : (K - 1) * s_h_t + V ≤ K * V) :
∃ sF, exec (cla_bwd_dh_surface q do_ dh s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t
s_vo_d s_h_h s_h_t scale T K V BT BK BV NT).toAlgKernel s = some sF
∧ ∀ t (idx : TileIndex [BK, BV]), t < NT → claActive s K V BK BV idx →
sF.readMem dh (claHOffset s s_h_h s_h_t K V BK BV t idx)
= claDhState s q do_ s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d scale
T K V BT BK BV NT t idx.1.val idx.2.1.val
Referenced model/specification definitions: Python source |
|---|
chunk_retentioncrh_fwd_h_exec_genuineMathematical execution | Configured model / stage Forward recurrent-state and backward-state stages with their particular layout and tile hypotheses. No output/dqkv stage composition or frame is present in these headlines. | Statement and preconditionsspecification crh_fwd_h_exec_genuine
(k v h h0 ht : RegionName)
(s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d s_h_h s_h_t : Nat)
(H T K V BT BK BV NT : Nat) (USE_INITIAL_STATE STORE_FINAL_STATE : Bool)
(s : BlockState)
(hHk : h ≠ k) (hHv : h ≠ v) (hHtH : ht ≠ h)
(hσ : BV ≤ s_h_t) (hFit : (K - 1) * s_h_t + V ≤ K * V) (hBVV : BV ≤ V) :
∃ sF, exec (crh_fwd_h_surface k v h h0 ht s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t
s_vo_d s_h_h s_h_t H T K V BT BK BV NT
USE_INITIAL_STATE STORE_FINAL_STATE).toAlgKernel s = some sF
∧ (∀ t (idx : TileIndex [BK, BV]), t < NT → crhActive s K V BK BV idx →
sF.readMem h (crhHOffset s s_h_h s_h_t K V BK BV t idx)
= crhState s k v h0 USE_INITIAL_STATE s_qk_h s_qk_t s_qk_d s_vo_h
s_vo_t s_vo_d H T K V BT BK BV NT t idx.1.val idx.2.1.val)
∧ (STORE_FINAL_STATE = Bool.true →
∀ idx : TileIndex [BK, BV], crhActive s K V BK BV idx →
sF.readMem ht (crhHtOffset s K V BK BV idx)
= crhState s k v h0 USE_INITIAL_STATE s_qk_h s_qk_t s_qk_d s_vo_h
s_vo_t s_vo_d H T K V BT BK BV NT NT idx.1.val idx.2.1.val)
Referenced model/specification definitions: Python source |
|---|
chunk_retentioncrh_bwd_dh_exec_genuineMathematical execution | Configured model / stage Forward recurrent-state and backward-state stages with their particular layout and tile hypotheses. No output/dqkv stage composition or frame is present in these headlines. | Statement and preconditionsspecification crh_bwd_dh_exec_genuine
(v do_ dh : RegionName)
(s_vo_h s_vo_t s_vo_d s_h_h s_h_t : Nat)
(H T K V BT NT : Nat) (s : BlockState)
(hσ : BT ≤ s_h_t) :
∃ sF, exec (crh_bwd_dh_surface v do_ dh s_vo_h s_vo_t s_vo_d s_h_h s_h_t
H T K V BT NT).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BT, BT], crhDhActive s K V BT idx →
sF.readMem dh (crhDhOffset s s_h_h s_h_t K V BT idx)
= crhDhOut s v do_ s_vo_h s_vo_t s_vo_d H T V BT NT
idx.1.val idx.2.1.val
Referenced model/specification definitions: Python source |
|---|
chunk_retention_opscro_fwd_h_exec_genuineMathematical execution | Configured model / stage Forward recurrent-state and backward dh stages only; other Python stages and host composition are outside these exec-existential contracts, which omit a frame. | Statement and preconditionsspecification cro_fwd_h_exec_genuine
(k v h h0 ht : RegionName)
(s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d s_h_h s_h_t : Nat)
(H T K V BT BK BV NT : Nat) (USE_INITIAL_STATE STORE_FINAL_STATE : Bool)
(s : BlockState)
(hHk : h ≠ k) (hHv : h ≠ v) (hHtH : ht ≠ h)
(hσ : BV ≤ s_h_t) (hFit : (K - 1) * s_h_t + V ≤ K * V) (hBVV : BV ≤ V) :
∃ sF, exec (cro_fwd_h_surface k v h h0 ht s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t
s_vo_d s_h_h s_h_t H T K V BT BK BV NT
USE_INITIAL_STATE STORE_FINAL_STATE).toAlgKernel s = some sF
∧ (∀ t (idx : TileIndex [BK, BV]), t < NT → croActive s K V BK BV idx →
sF.readMem h (croHOffset s s_h_h s_h_t K V BK BV t idx)
= croState s k v h0 USE_INITIAL_STATE s_qk_h s_qk_t s_qk_d s_vo_h
s_vo_t s_vo_d H T K V BT BK BV NT t idx.1.val idx.2.1.val)
∧ (STORE_FINAL_STATE = Bool.true →
∀ idx : TileIndex [BK, BV], croActive s K V BK BV idx →
sF.readMem ht (croHtOffset s K V BK BV idx)
= croState s k v h0 USE_INITIAL_STATE s_qk_h s_qk_t s_qk_d s_vo_h
s_vo_t s_vo_d H T K V BT BK BV NT NT idx.1.val idx.2.1.val)
Referenced model/specification definitions: Python source |
|---|
chunk_retention_opscro_bwd_dh_exec_genuineMathematical execution | Configured model / stage Forward recurrent-state and backward dh stages only; other Python stages and host composition are outside these exec-existential contracts, which omit a frame. | Statement and preconditionsspecification cro_bwd_dh_exec_genuine
(q do_ dh : RegionName)
(s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d s_h_h s_h_t : Nat)
(scale : ℝ) (H T K V BT BK BV NT : Nat) (s : BlockState)
(hDq : dh ≠ q) (hDdo : dh ≠ do_)
(hσ : BV ≤ s_h_t) (hFit : (K - 1) * s_h_t + V ≤ K * V) :
∃ sF, exec (cro_bwd_dh_surface q do_ dh s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t
s_vo_d s_h_h s_h_t scale H T K V BT BK BV NT).toAlgKernel s = some sF
∧ ∀ t (idx : TileIndex [BK, BV]), t < NT → croActive s K V BK BV idx →
sF.readMem dh (croHOffset s s_h_h s_h_t K V BK BV t idx)
= croDhCarry s q do_ s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d
scale H T K V BT BK BV NT (NT - 1 - t) idx.1.val idx.2.1.val
Referenced model/specification definitions: Python source |
|---|
chunked_cumsum_fwdchunked_cumsum_fwd_correctnessMathematical execution | Precomputed-input slice The IO contract consumes DtPrepared and proves dt storage plus dA cumulative sums. The combined summary also contains full-surface facts; those do not make the prepared-input IO kernel the original preprocessing kernel. | Statement and preconditionsspecification chunked_cumsum_fwd_correctness
(DtPrepared A DtOut DACumsum : RegionName)
(stride_dt_batch stride_dt_seqlen stride_dt_head stride_A_head
stride_dt_out_batch stride_dt_out_chunk stride_dt_out_head stride_dt_out_csize
stride_dA_cs_batch stride_dA_cs_chunk stride_dA_cs_head stride_dA_cs_csize
nheads chunk_size BLOCK_SIZE_H BLOCK_SIZE_CHUNK : Nat)
(hDtNe : DtOut ≠ DACumsum) (hCK : 0 < BLOCK_SIZE_CHUNK)
(hDtOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_SIZE_H, BLOCK_SIZE_CHUNK] =>
p₀ * stride_dt_out_batch + p₁ * stride_dt_out_chunk +
(p₂ * BLOCK_SIZE_H + idx.1.val) * stride_dt_out_head +
idx.2.1.val * stride_dt_out_csize))
(hDACsInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_SIZE_H, BLOCK_SIZE_CHUNK] =>
p₀ * stride_dA_cs_batch + p₁ * stride_dA_cs_chunk +
(p₂ * BLOCK_SIZE_H + idx.1.val) * stride_dA_cs_head +
idx.2.1.val * stride_dA_cs_csize)) :
chunkedCumsumFwdIO DtPrepared A DtOut DACumsum
stride_dt_batch stride_dt_seqlen stride_dt_head stride_A_head
stride_dt_out_batch stride_dt_out_chunk stride_dt_out_head stride_dt_out_csize
stride_dA_cs_batch stride_dA_cs_chunk stride_dA_cs_head stride_dA_cs_csize
nheads chunk_size BLOCK_SIZE_H BLOCK_SIZE_CHUNK ⊨
fun _ _ xs ys =>
(fun j => xs j,
fun j => dACumsumLaneClosed chunk_size BLOCK_SIZE_H BLOCK_SIZE_CHUNK xs ys j)
IO definition: chunkedCumsumFwdIOdef chunkedCumsumFwdIO (DtPrepared A DtOut DACumsum : RegionName)
(stride_dt_batch stride_dt_seqlen stride_dt_head stride_A_head
stride_dt_out_batch stride_dt_out_chunk stride_dt_out_head stride_dt_out_csize
stride_dA_cs_batch stride_dA_cs_chunk stride_dA_cs_head stride_dA_cs_csize
nheads chunk_size BLOCK_SIZE_H BLOCK_SIZE_CHUNK : Nat) :
Masked3DKernelIO₂ₓ₂ where
kernel := chunked_cumsum_fused_slice DtPrepared A DtOut DACumsum
stride_dt_batch stride_dt_seqlen stride_dt_head stride_A_head
stride_dt_out_batch stride_dt_out_chunk stride_dt_out_head stride_dt_out_csize
stride_dA_cs_batch stride_dA_cs_chunk stride_dA_cs_head stride_dA_cs_csize
nheads chunk_size BLOCK_SIZE_H BLOCK_SIZE_CHUNK
in1 := DtPrepared
in2 := A
out1 := DtOut
out2 := DACumsum
B := BLOCK_SIZE_H * BLOCK_SIZE_CHUNK
read1 := fun p₀ p₁ p₂ j =>
p₀ * stride_dt_batch +
(p₁ * chunk_size + j.val % BLOCK_SIZE_CHUNK) * stride_dt_seqlen +
(p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK) * stride_dt_head
read2 := fun _ _ p₂ j =>
(p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK) * stride_A_head
write1 := fun p₀ p₁ p₂ j =>
p₀ * stride_dt_out_batch + p₁ * stride_dt_out_chunk +
(p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK) * stride_dt_out_head +
j.val % BLOCK_SIZE_CHUNK * stride_dt_out_csize
write2 := fun p₀ p₁ p₂ j =>
p₀ * stride_dA_cs_batch + p₁ * stride_dA_cs_chunk +
(p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK) * stride_dA_cs_head +
j.val % BLOCK_SIZE_CHUNK * stride_dA_cs_csize
mask := fun _ _ p₂ j =>
p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK < nheads ∧
j.val % BLOCK_SIZE_CHUNK < chunk_size
read2Mask := fun _ _ p₂ j => p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK < nheads
Referenced model/specification definitions: Python source |
|---|
chunked_cumsum_fwdchunked_cumsum_fwd_summary_generalMathematical execution | Configured model / stage A conjunction of full-surface projection/value facts and a distinct DtPrepared IO slice; inspect each conjunct rather than treating the prepared-input skin as the entire Python kernel. | Statement and preconditionsspecification chunked_cumsum_fwd_summary_general
(dt_ptr A_ptr dt_bias_ptr dt_out_ptr dA_cumsum_ptr
DtPrepared A DtOut DACumsum : RegionName)
(batch seqlen nheads chunk_size : Nat)
(dt_min dt_max : ℝ)
(stride_dt_batch stride_dt_seqlen stride_dt_head stride_A_head
stride_dt_bias_head
stride_dt_out_batch stride_dt_out_chunk stride_dt_out_head stride_dt_out_csize
stride_dA_cs_batch stride_dA_cs_chunk stride_dA_cs_head stride_dA_cs_csize : Nat)
(DT_SOFTPLUS HAS_DT_BIAS : Bool)
(BLOCK_SIZE_H BLOCK_SIZE_CHUNK : Nat)
(hDtNe : DtOut ≠ DACumsum) (hCK : 0 < BLOCK_SIZE_CHUNK)
(hDtOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_SIZE_H, BLOCK_SIZE_CHUNK] =>
p₀ * stride_dt_out_batch + p₁ * stride_dt_out_chunk +
(p₂ * BLOCK_SIZE_H + idx.1.val) * stride_dt_out_head +
idx.2.1.val * stride_dt_out_csize))
(hDACsInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_SIZE_H, BLOCK_SIZE_CHUNK] =>
p₀ * stride_dA_cs_batch + p₁ * stride_dA_cs_chunk +
(p₂ * BLOCK_SIZE_H + idx.1.val) * stride_dA_cs_head +
idx.2.1.val * stride_dA_cs_csize)) :
(∃ alg, (chunked_cumsum_fwd_surface dt_ptr A_ptr dt_bias_ptr dt_out_ptr
dA_cumsum_ptr batch seqlen nheads chunk_size dt_min dt_max stride_dt_batch
stride_dt_seqlen stride_dt_head stride_A_head stride_dt_bias_head
stride_dt_out_batch stride_dt_out_chunk stride_dt_out_head stride_dt_out_csize
stride_dA_cs_batch stride_dA_cs_chunk stride_dA_cs_head stride_dA_cs_csize
DT_SOFTPLUS HAS_DT_BIAS BLOCK_SIZE_H BLOCK_SIZE_CHUNK).toAlgorithm?
= Except.ok alg) ∧
(chunkedCumsumFwdIO DtPrepared A DtOut DACumsum
stride_dt_batch stride_dt_seqlen stride_dt_head stride_A_head
stride_dt_out_batch stride_dt_out_chunk stride_dt_out_head stride_dt_out_csize
stride_dA_cs_batch stride_dA_cs_chunk stride_dA_cs_head stride_dA_cs_csize
nheads chunk_size BLOCK_SIZE_H BLOCK_SIZE_CHUNK ⊨
fun _ _ xs ys =>
(fun j => xs j,
fun j => dACumsumLaneClosed chunk_size BLOCK_SIZE_H BLOCK_SIZE_CHUNK xs ys j))
IO definition: chunkedCumsumFwdIOdef chunkedCumsumFwdIO (DtPrepared A DtOut DACumsum : RegionName)
(stride_dt_batch stride_dt_seqlen stride_dt_head stride_A_head
stride_dt_out_batch stride_dt_out_chunk stride_dt_out_head stride_dt_out_csize
stride_dA_cs_batch stride_dA_cs_chunk stride_dA_cs_head stride_dA_cs_csize
nheads chunk_size BLOCK_SIZE_H BLOCK_SIZE_CHUNK : Nat) :
Masked3DKernelIO₂ₓ₂ where
kernel := chunked_cumsum_fused_slice DtPrepared A DtOut DACumsum
stride_dt_batch stride_dt_seqlen stride_dt_head stride_A_head
stride_dt_out_batch stride_dt_out_chunk stride_dt_out_head stride_dt_out_csize
stride_dA_cs_batch stride_dA_cs_chunk stride_dA_cs_head stride_dA_cs_csize
nheads chunk_size BLOCK_SIZE_H BLOCK_SIZE_CHUNK
in1 := DtPrepared
in2 := A
out1 := DtOut
out2 := DACumsum
B := BLOCK_SIZE_H * BLOCK_SIZE_CHUNK
read1 := fun p₀ p₁ p₂ j =>
p₀ * stride_dt_batch +
(p₁ * chunk_size + j.val % BLOCK_SIZE_CHUNK) * stride_dt_seqlen +
(p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK) * stride_dt_head
read2 := fun _ _ p₂ j =>
(p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK) * stride_A_head
write1 := fun p₀ p₁ p₂ j =>
p₀ * stride_dt_out_batch + p₁ * stride_dt_out_chunk +
(p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK) * stride_dt_out_head +
j.val % BLOCK_SIZE_CHUNK * stride_dt_out_csize
write2 := fun p₀ p₁ p₂ j =>
p₀ * stride_dA_cs_batch + p₁ * stride_dA_cs_chunk +
(p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK) * stride_dA_cs_head +
j.val % BLOCK_SIZE_CHUNK * stride_dA_cs_csize
mask := fun _ _ p₂ j =>
p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK < nheads ∧
j.val % BLOCK_SIZE_CHUNK < chunk_size
read2Mask := fun _ _ p₂ j => p₂ * BLOCK_SIZE_H + j.val / BLOCK_SIZE_CHUNK < nheads
Referenced model/specification definitions: Python source |
|---|
context_attn_bloomcontext_attn_bloom_surface_compute_correct_generalMathematical execution | Configured model / stage Prompt-cache context attention with explicit natural metadata, token mapping, head dimensions, and stream extent. The IO binding fixes launch/layout arguments shown in its definition. | Statement and preconditionsspecification context_attn_bloom_surface_compute_correct_general
(Q K V Out B_Start_Loc B_Seqlen Req_to_tokens B_req_idx B_Prompt_Cache_Len : RegionName)
(sm_scale : ℝ)
(stride_qbs stride_qh stride_qd stride_req_b stride_req_s stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
head_dim BLOCK_DMODEL BLOCK_M BLOCK_N : Nat) (hD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N)
(s : BlockState)
(hOInj : Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => outOffset s B_Start_Loc stride_obs stride_oh stride_od BLOCK_M idx))
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := context_attn_bloom_fwd_kernel_surface Q K V sm_scale
B_Start_Loc B_Seqlen Out Req_to_tokens B_req_idx B_Prompt_Cache_Len
stride_qbs stride_qh stride_qd stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
stride_req_b stride_req_s 1 head_dim BLOCK_M BLOCK_DMODEL BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
active s B_Seqlen B_Prompt_Cache_Len head_dim BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(Out, outOffset s B_Start_Loc stride_obs stride_oh stride_od BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
bloomFwdGenuineOutValueG s Q K V B_Start_Loc B_Seqlen Req_to_tokens B_req_idx B_Prompt_Cache_Len
sm_scale stride_qbs stride_qh stride_qd stride_req_b stride_req_s stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd head_dim BLOCK_DMODEL BLOCK_N BLOCK_M idx)
Referenced model/specification definitions: Python source |
|---|
context_attn_bloomcontext_attn_bloom_io_correctnessAbstract cast/store rounding | Configured model / stage Prompt-cache context attention with explicit natural metadata, token mapping, head dimensions, and stream extent. The IO binding fixes launch/layout arguments shown in its definition. | Statement and preconditionsspecification context_attn_bloom_io_correctness (R : RoundingModel)
(Q K V Out : RegionName)
(B_Start_Loc B_Seqlen Req_to_tokens B_req_idx b_prompt_cache_len : Region .nat)
(sm_scale : ℝ)
(stride_qbs stride_qh stride_qd stride_req_b stride_req_s stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
head_dim BLOCK_DMODEL BLOCK_M BLOCK_N NT : Nat)
(hD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M) (hNT : 0 < NT)
(hOInj : ∀ pid₁ pid₂ base : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(base + (pid₂ * BLOCK_M + idx.1.val)) * stride_obs + pid₁ * stride_oh
+ idx.2.1.val * stride_od)) :
contextAttnBloomIO Q K V Out B_Start_Loc B_Seqlen Req_to_tokens B_req_idx
b_prompt_cache_len sm_scale
stride_qbs stride_qh stride_qd stride_req_b stride_req_s stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
head_dim BLOCK_DMODEL BLOCK_M BLOCK_N NT ⊨[R]
fun _ _ pid₂ m xs ys zs j =>
contextAttnBloomIOSpec BLOCK_M BLOCK_DMODEL BLOCK_N head_dim NT hBN
(bloomFwdIOT_pos NT BLOCK_M BLOCK_N hBM hBN hNT) sm_scale pid₂
(m (⟨1, by omega⟩ : Fin 4)) (m (⟨2, by omega⟩ : Fin 4)) xs ys zs j
IO definition: contextAttnBloomIOdef contextAttnBloomIO (Q K V Out : RegionName)
(B_Start_Loc B_Seqlen Req_to_tokens B_req_idx b_prompt_cache_len : Region .nat) (sm_scale : ℝ)
(stride_qbs stride_qh stride_qd stride_req_b stride_req_s stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
head_dim BLOCK_DMODEL BLOCK_M BLOCK_N NT : Nat) :
StreamMetaGatherMasked3DKernelIO₃ where
kernel := context_attn_bloom_fwd_kernel_surface Q K V sm_scale B_Start_Loc B_Seqlen Out
Req_to_tokens B_req_idx b_prompt_cache_len
stride_qbs stride_qh stride_qd stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
stride_req_b stride_req_s 1 head_dim BLOCK_M BLOCK_DMODEL BLOCK_N
inp1 := Q
inp2 := K
inp3 := V
out := Out
nMeta := 4
sty := fun _ => ChanTy.nat
mbuf := bloomIOMetaBuf B_Start_Loc B_Seqlen B_req_idx b_prompt_cache_len
mwin := fun _ pid₀ _ _ => pid₀
gbuf := Req_to_tokens
gty := ChanTy.nat
Bg := BLOCK_N
gother := 0
T := bloomFwdIOT NT BLOCK_M BLOCK_N
B1 := BLOCK_M * BLOCK_DMODEL
B2 := BLOCK_DMODEL * BLOCK_N
B3 := BLOCK_N * BLOCK_DMODEL
C := BLOCK_M * BLOCK_DMODEL
pre := fun _ _ pid₂ m =>
pid₂ < NT ∧ m (⟨2, by omega⟩ : Fin 4) ≤ NT * BLOCK_M
gread := fun _ _ _ m t j =>
stride_req_b * m (⟨3, by omega⟩ : Fin 4) + stride_req_s * (t.val * BLOCK_N + j.val)
gmask := fun _ _ pid₂ m t j =>
t.val * BLOCK_N + j.val
< bloomFwdIOBel BLOCK_M pid₂ (m (⟨1, by omega⟩ : Fin 4)) (m (⟨2, by omega⟩ : Fin 4))
read1 := fun _ pid₁ pid₂ m _ j =>
(m (⟨0, by omega⟩ : Fin 4) + (pid₂ * BLOCK_M + j.val / BLOCK_DMODEL)) * stride_qbs
+ pid₁ * stride_qh + j.val % BLOCK_DMODEL * stride_qd
read2 := fun _ pid₁ _ _ G t j =>
G t (Lane2D.decode j).2.1 * stride_kbs
+ pid₁ * stride_kh + j.val / BLOCK_N * stride_kd
read3 := fun _ pid₁ _ _ G t j =>
G t (Lane2D.decode j).1 * stride_vbs
+ pid₁ * stride_vh + j.val % BLOCK_DMODEL * stride_vd
write := fun _ pid₁ pid₂ m j =>
(m (⟨0, by omega⟩ : Fin 4) + (pid₂ * BLOCK_M + j.val / BLOCK_DMODEL)) * stride_obs
+ pid₁ * stride_oh + j.val % BLOCK_DMODEL * stride_od
mask1 := fun _ _ pid₂ m _ j =>
pid₂ * BLOCK_M + j.val / BLOCK_DMODEL
< m (⟨2, by omega⟩ : Fin 4) - m (⟨1, by omega⟩ : Fin 4)
∧ j.val % BLOCK_DMODEL < head_dim
mask2 := fun _ _ pid₂ m t j =>
t.val * BLOCK_N + j.val % BLOCK_N
< bloomFwdIOBel BLOCK_M pid₂ (m (⟨1, by omega⟩ : Fin 4)) (m (⟨2, by omega⟩ : Fin 4))
∧ j.val / BLOCK_N < head_dim
mask3 := fun _ _ pid₂ m t j =>
t.val * BLOCK_N + j.val / BLOCK_DMODEL
< bloomFwdIOBel BLOCK_M pid₂ (m (⟨1, by omega⟩ : Fin 4)) (m (⟨2, by omega⟩ : Fin 4))
∧ j.val % BLOCK_DMODEL < head_dim
writeMask := fun _ _ pid₂ m j =>
pid₂ * BLOCK_M + j.val / BLOCK_DMODEL
< m (⟨2, by omega⟩ : Fin 4) - m (⟨1, by omega⟩ : Fin 4)
∧ j.val % BLOCK_DMODEL < head_dim
Referenced model/specification definitions: Python source |
|---|
context_attn_fwdcontext_attn_fwd_surface_compute_correct_generalMathematical execution | Configured model / stage Int8-KV context-attention model with explicit metadata and stride assumptions; mathematical casts are not a proof of fixed-width hardware quantization. | Statement and preconditionsspecification context_attn_fwd_surface_compute_correct_general
(Q K V Out B_Start_Loc B_Seqlen B_Prompt_Cache_Len : RegionName)
(sm_scale : ℝ)
(stride_qbs stride_qh stride_qd stride_kb stride_kh stride_ks stride_kd
stride_vb stride_vh stride_vs stride_vd stride_obs stride_oh stride_od
H BLOCK_DMODEL BLOCK_M BLOCK_N : Nat) (hD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N)
(s : BlockState)
(hOInj : Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => outOffset s H B_Start_Loc stride_obs stride_oh stride_od BLOCK_M idx))
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := context_attn_fwd_kernel_int8kv_surface Q K V sm_scale Out
B_Start_Loc B_Seqlen B_Prompt_Cache_Len
stride_qbs stride_qh stride_qd stride_kb stride_kh stride_ks stride_kd
stride_vb stride_vh stride_vs stride_vd stride_obs stride_oh stride_od
1 H BLOCK_DMODEL BLOCK_M BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
active s H B_Seqlen B_Prompt_Cache_Len BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(Out, outOffset s H B_Start_Loc stride_obs stride_oh stride_od BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
ctxFwdGenuineOutValueG s Q K V B_Start_Loc B_Seqlen B_Prompt_Cache_Len
sm_scale H stride_qbs stride_qh stride_qd stride_kb stride_ks stride_kh stride_kd
stride_vb stride_vs stride_vh stride_vd BLOCK_DMODEL BLOCK_M BLOCK_N idx)
Referenced model/specification definitions: Python source |
|---|
context_attn_fwdcontext_attn_fwd_io_correctnessAbstract cast/store rounding | Configured model / stage Int8-KV context-attention model with explicit metadata and stride assumptions; mathematical casts are not a proof of fixed-width hardware quantization. | Statement and preconditionsspecification context_attn_fwd_io_correctness (R : RoundingModel)
(Q K V Out : RegionName)
(B_Start_Loc B_Seqlen b_prompt_cache_len : Region .nat)
(sm_scale : ℝ)
(stride_qbs stride_qh stride_qd stride_kb stride_kh stride_ks stride_kd
stride_vb stride_vh stride_vs stride_vd stride_obs stride_oh stride_od
H BLOCK_DMODEL BLOCK_M BLOCK_N NT : Nat)
(hD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M) (hNT : 0 < NT)
(hOInj : ∀ pid₀ pid₁ base : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(base + (pid₀ * BLOCK_M + idx.1.val)) * stride_obs + pid₁ % H * stride_oh
+ idx.2.1.val * stride_od)) :
contextAttnFwdIO Q K V Out B_Start_Loc B_Seqlen b_prompt_cache_len sm_scale
stride_qbs stride_qh stride_qd stride_kb stride_kh stride_ks stride_kd
stride_vb stride_vh stride_vs stride_vd stride_obs stride_oh stride_od
H BLOCK_DMODEL BLOCK_M BLOCK_N NT ⊨[R]
fun pid₀ _ _ m xs ys zs j =>
contextAttnFwdIOSpec BLOCK_M BLOCK_DMODEL BLOCK_N NT hBN
(ctxFwdIOT_pos NT BLOCK_M BLOCK_N hBM hBN hNT) sm_scale pid₀
(m (⟨0, by omega⟩ : Fin 3)) (m (⟨2, by omega⟩ : Fin 3)) xs ys zs j
IO definition: contextAttnFwdIOdef contextAttnFwdIO (Q K V Out : RegionName)
(B_Start_Loc B_Seqlen b_prompt_cache_len : Region .nat) (sm_scale : ℝ)
(stride_qbs stride_qh stride_qd stride_kb stride_kh stride_ks stride_kd
stride_vb stride_vh stride_vs stride_vd stride_obs stride_oh stride_od
H BLOCK_DMODEL BLOCK_M BLOCK_N NT : Nat) :
StreamMetaMasked3DKernelIO₃ where
kernel := context_attn_fwd_kernel_int8kv_surface Q K V sm_scale Out
B_Start_Loc B_Seqlen b_prompt_cache_len
stride_qbs stride_qh stride_qd stride_kb stride_kh stride_ks stride_kd
stride_vb stride_vh stride_vs stride_vd stride_obs stride_oh stride_od
1 H BLOCK_DMODEL BLOCK_M BLOCK_N
inp1 := Q
inp2 := K
inp3 := V
out := Out
nMeta := 3
sty := fun _ => ChanTy.nat
mbuf := ctxFwdIOMetaBuf B_Start_Loc B_Seqlen b_prompt_cache_len
mwin := fun _ _ pid₁ _ => pid₁ / H
T := ctxFwdIOT NT BLOCK_M BLOCK_N
B1 := BLOCK_M * BLOCK_DMODEL
B2 := BLOCK_DMODEL * BLOCK_N
B3 := BLOCK_N * BLOCK_DMODEL
C := BLOCK_M * BLOCK_DMODEL
pre := fun pid₀ _ _ m =>
pid₀ < NT ∧ m (⟨2, by omega⟩ : Fin 3) ≤ NT * BLOCK_M
read1 := fun pid₀ pid₁ _ m _ j =>
(m (⟨1, by omega⟩ : Fin 3) + (pid₀ * BLOCK_M + j.val / BLOCK_DMODEL)) * stride_qbs
+ pid₁ % H * stride_qh + j.val % BLOCK_DMODEL * stride_qd
read2 := fun _ pid₁ _ _ t j =>
pid₁ / H * stride_kb + (t.val * BLOCK_N + j.val % BLOCK_N) * stride_ks
+ pid₁ % H * stride_kh + j.val / BLOCK_N * stride_kd
read3 := fun _ pid₁ _ _ t j =>
pid₁ / H * stride_vb + (t.val * BLOCK_N + j.val / BLOCK_DMODEL) * stride_vs
+ pid₁ % H * stride_vh + j.val % BLOCK_DMODEL * stride_vd
write := fun pid₀ pid₁ _ m j =>
(m (⟨1, by omega⟩ : Fin 3) + (pid₀ * BLOCK_M + j.val / BLOCK_DMODEL)) * stride_obs
+ pid₁ % H * stride_oh + j.val % BLOCK_DMODEL * stride_od
mask1 := fun pid₀ _ _ m _ j =>
pid₀ * BLOCK_M + j.val / BLOCK_DMODEL
< m (⟨2, by omega⟩ : Fin 3) - m (⟨0, by omega⟩ : Fin 3)
mask2 := fun pid₀ _ _ m t j =>
t.val * BLOCK_N + j.val % BLOCK_N
< ctxFwdIOBel BLOCK_M pid₀ (m (⟨0, by omega⟩ : Fin 3)) (m (⟨2, by omega⟩ : Fin 3))
mask3 := fun pid₀ _ _ m t j =>
t.val * BLOCK_N + j.val / BLOCK_DMODEL
< ctxFwdIOBel BLOCK_M pid₀ (m (⟨0, by omega⟩ : Fin 3)) (m (⟨2, by omega⟩ : Fin 3))
writeMask := fun pid₀ _ _ m j =>
pid₀ * BLOCK_M + j.val / BLOCK_DMODEL
< m (⟨2, by omega⟩ : Fin 3) - m (⟨0, by omega⟩ : Fin 3)
Referenced model/specification definitions: Python source |
|---|
context_attn_llamacontext_attn_llama_surface_compute_correct_generalMathematical execution | Configured model / stage LLaMA prompt-cache context attention with supplied token mapping and sequence metadata; IO stream extent, bounds, and layout assumptions are explicit. | Statement and preconditionsspecification context_attn_llama_surface_compute_correct_general
(Q K V Out B_Start_Loc B_Seqlen Req_to_tokens B_req_idx B_Prompt_Cache_Len : RegionName)
(sm_scale : ℝ)
(stride_qbs stride_qh stride_qd stride_req_b stride_req_s stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
H BLOCK_DMODEL BLOCK_M BLOCK_N : Nat) (hD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N)
(s : BlockState)
(hOInj : Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] => outOffset s H B_Start_Loc stride_obs stride_oh stride_od BLOCK_M idx))
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := context_attn_llama_fwd_kernel_surface Q K V sm_scale Out
B_Start_Loc B_Seqlen Req_to_tokens B_req_idx B_Prompt_Cache_Len
stride_qbs stride_qh stride_qd stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
stride_req_b stride_req_s 1 H BLOCK_DMODEL BLOCK_M BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
active s H B_Seqlen B_Prompt_Cache_Len BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(Out, outOffset s H B_Start_Loc stride_obs stride_oh stride_od BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
ctxFwdGenuineOutValueG s Q K V B_Start_Loc B_Seqlen Req_to_tokens B_req_idx B_Prompt_Cache_Len
sm_scale H stride_qbs stride_qh stride_qd stride_req_b stride_req_s stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd BLOCK_DMODEL BLOCK_M BLOCK_N idx)
Referenced model/specification definitions: Python source |
|---|
context_attn_llamacontext_attn_llama_io_correctnessAbstract cast/store rounding | Configured model / stage LLaMA prompt-cache context attention with supplied token mapping and sequence metadata; IO stream extent, bounds, and layout assumptions are explicit. | Statement and preconditionsspecification context_attn_llama_io_correctness (R : RoundingModel)
(Q K V Out : RegionName)
(B_Start_Loc B_Seqlen Req_to_tokens B_req_idx b_prompt_cache_len : Region .nat)
(sm_scale : ℝ)
(stride_qbs stride_qh stride_qd stride_req_b stride_req_s stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
H BLOCK_DMODEL BLOCK_M BLOCK_N NT : Nat)
(hD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M) (hNT : 0 < NT)
(hOInj : ∀ pid₀ pid₁ base : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(base + (pid₀ * BLOCK_M + idx.1.val)) * stride_obs + pid₁ % H * stride_oh
+ idx.2.1.val * stride_od)) :
contextAttnLlamaIO Q K V Out B_Start_Loc B_Seqlen Req_to_tokens B_req_idx
b_prompt_cache_len sm_scale
stride_qbs stride_qh stride_qd stride_req_b stride_req_s stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
H BLOCK_DMODEL BLOCK_M BLOCK_N NT ⊨[R]
fun pid₀ _ _ m xs ys zs j =>
contextAttnLlamaIOSpec BLOCK_M BLOCK_DMODEL BLOCK_N NT hBN
(ctxFwdIOT_pos NT BLOCK_M BLOCK_N hBM hBN hNT) sm_scale pid₀
(m (⟨1, by omega⟩ : Fin 4)) (m (⟨2, by omega⟩ : Fin 4)) xs ys zs j
IO definition: contextAttnLlamaIOdef contextAttnLlamaIO (Q K V Out : RegionName)
(B_Start_Loc B_Seqlen Req_to_tokens B_req_idx b_prompt_cache_len : Region .nat) (sm_scale : ℝ)
(stride_qbs stride_qh stride_qd stride_req_b stride_req_s stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
H BLOCK_DMODEL BLOCK_M BLOCK_N NT : Nat) :
StreamMetaGatherMasked3DKernelIO₃ where
kernel := context_attn_llama_fwd_kernel_surface Q K V sm_scale Out
B_Start_Loc B_Seqlen Req_to_tokens B_req_idx b_prompt_cache_len
stride_qbs stride_qh stride_qd stride_kbs stride_kh stride_kd
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
stride_req_b stride_req_s 1 H BLOCK_DMODEL BLOCK_M BLOCK_N
inp1 := Q
inp2 := K
inp3 := V
out := Out
nMeta := 4
sty := fun _ => ChanTy.nat
mbuf := ctxLlamaIOMetaBuf B_Start_Loc B_Seqlen B_req_idx b_prompt_cache_len
mwin := fun _ _ pid₁ _ => pid₁ / H
gbuf := Req_to_tokens
gty := ChanTy.nat
Bg := BLOCK_N
gother := 0
T := ctxFwdIOT NT BLOCK_M BLOCK_N
B1 := BLOCK_M * BLOCK_DMODEL
B2 := BLOCK_DMODEL * BLOCK_N
B3 := BLOCK_N * BLOCK_DMODEL
C := BLOCK_M * BLOCK_DMODEL
pre := fun pid₀ _ _ m =>
pid₀ < NT ∧ m (⟨2, by omega⟩ : Fin 4) ≤ NT * BLOCK_M
gread := fun _ _ _ m t j =>
stride_req_b * m (⟨3, by omega⟩ : Fin 4) + stride_req_s * (t.val * BLOCK_N + j.val)
gmask := fun pid₀ _ _ m t j =>
t.val * BLOCK_N + j.val
< ctxFwdIOBel BLOCK_M pid₀ (m (⟨1, by omega⟩ : Fin 4)) (m (⟨2, by omega⟩ : Fin 4))
read1 := fun pid₀ pid₁ _ m _ j =>
(m (⟨0, by omega⟩ : Fin 4) + (pid₀ * BLOCK_M + j.val / BLOCK_DMODEL)) * stride_qbs
+ pid₁ % H * stride_qh + j.val % BLOCK_DMODEL * stride_qd
read2 := fun _ pid₁ _ _ G t j =>
G t (Lane2D.decode j).2.1 * stride_kbs
+ pid₁ % H * stride_kh + j.val / BLOCK_N * stride_kd
read3 := fun _ pid₁ _ _ G t j =>
G t (Lane2D.decode j).1 * stride_vbs
+ pid₁ % H * stride_vh + j.val % BLOCK_DMODEL * stride_vd
write := fun pid₀ pid₁ _ m j =>
(m (⟨0, by omega⟩ : Fin 4) + (pid₀ * BLOCK_M + j.val / BLOCK_DMODEL)) * stride_obs
+ pid₁ % H * stride_oh + j.val % BLOCK_DMODEL * stride_od
mask1 := fun pid₀ _ _ m _ j =>
pid₀ * BLOCK_M + j.val / BLOCK_DMODEL
< m (⟨2, by omega⟩ : Fin 4) - m (⟨1, by omega⟩ : Fin 4)
mask2 := fun pid₀ _ _ m t j =>
t.val * BLOCK_N + j.val % BLOCK_N
< ctxFwdIOBel BLOCK_M pid₀ (m (⟨1, by omega⟩ : Fin 4)) (m (⟨2, by omega⟩ : Fin 4))
mask3 := fun pid₀ _ _ m t j =>
t.val * BLOCK_N + j.val / BLOCK_DMODEL
< ctxFwdIOBel BLOCK_M pid₀ (m (⟨1, by omega⟩ : Fin 4)) (m (⟨2, by omega⟩ : Fin 4))
writeMask := fun pid₀ _ _ m j =>
pid₀ * BLOCK_M + j.val / BLOCK_DMODEL
< m (⟨2, by omega⟩ : Fin 4) - m (⟨1, by omega⟩ : Fin 4)
Referenced model/specification definitions: Python source |
|---|
context_attn_mistralcontext_attn_mistral_genuine_output_summary_generalMathematical execution | Configured model / stage Sliding-window context attention with equal BLK query/key block sizes and shared contiguous layouts. Sequence metadata and positive stream-size premises remain required. | Statement and preconditionsspecification context_attn_mistral_genuine_output_summary_general
(Q K V : RegionName) (B_Start_Loc B_Seqlen : Region .nat)
(Out : RegionName) (sm_scale : ℝ) (rs hs BLK DM sw : Nat)
(hBLK : 0 < BLK) (hDM : 0 < DM) (hDMrs : DM ≤ rs)
(s : BlockState) (hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := context_attn_mistral_fwd_kernel_surface Q K V sm_scale B_Start_Loc B_Seqlen Out
rs hs 1 rs hs 1 rs hs 1 rs hs 1 1 sw BLK DM BLK)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLK, DM] => mistralActiveG s B_Seqlen BLK DM idx)
(fun idx : TileIndex [BLK, DM] => (Out, mistralOutOffsetG s B_Start_Loc rs hs BLK DM idx)))
(expected := fun idx : TileIndex [BLK, DM] =>
mistralGenuineOutValueG s Q K V B_Start_Loc B_Seqlen sm_scale rs hs BLK DM sw idx)
Referenced model/specification definitions: Python source |
|---|
context_attn_mistralcontext_attn_mistral_io_correctnessAbstract cast/store rounding | Configured model / stage Sliding-window context attention with equal BLK query/key block sizes and shared contiguous layouts. Sequence metadata and positive stream-size premises remain required. | Statement and preconditionsspecification context_attn_mistral_io_correctness (R : RoundingModel)
(Q K V : RegionName) (B_Start_Loc B_Seqlen : Region .nat) (Out : RegionName)
(sm_scale : ℝ) (rs hs sw BLK DM NT : Nat)
(hBLK : 0 < BLK) (hDM : 0 < DM) (hDMrs : DM ≤ rs) (hNT : 0 < NT) :
contextAttnMistralIO Q K V B_Start_Loc B_Seqlen Out sm_scale rs hs sw BLK DM NT ⊨[R]
fun _ _ pid₂ m xs ys zs j =>
contextAttnMistralIOSpec BLK DM NT sw (m (⟨0, by omega⟩ : Fin 2)) hBLK hNT sm_scale
pid₂ xs ys zs j
IO definition: contextAttnMistralIOdef contextAttnMistralIO (Q K V : RegionName) (B_Start_Loc B_Seqlen : Region .nat)
(Out : RegionName) (sm_scale : ℝ) (rs hs sw BLK DM NT : Nat) :
StreamMetaMasked3DKernelIO₃ where
kernel := context_attn_mistral_fwd_kernel_surface Q K V sm_scale B_Start_Loc B_Seqlen Out
rs hs 1 rs hs 1 rs hs 1 rs hs 1 1 sw BLK DM BLK
inp1 := Q
inp2 := K
inp3 := V
out := Out
nMeta := 2
sty := fun _ => ChanTy.nat
mbuf := ctxMistralMetaBuf B_Start_Loc B_Seqlen
mwin := fun _ pid₀ _ _ => pid₀
T := NT
B1 := BLK * DM
B2 := DM * BLK
B3 := BLK * DM
C := BLK * DM
pre := fun _ _ pid₂ m => pid₂ < NT ∧ m (⟨0, by omega⟩ : Fin 2) ≤ NT * BLK
read1 := fun _ pid₁ pid₂ m _ j =>
(m (⟨1, by omega⟩ : Fin 2) + (pid₂ * BLK + j.val / DM)) * rs + pid₁ * hs + j.val % DM
read2 := fun _ pid₁ _ m t j =>
(m (⟨1, by omega⟩ : Fin 2) + (t.val * BLK + j.val % BLK)) * rs + pid₁ * hs + j.val / BLK
read3 := fun _ pid₁ _ m t j =>
(m (⟨1, by omega⟩ : Fin 2) + (t.val * BLK + j.val / DM)) * rs + pid₁ * hs + j.val % DM
write := fun _ pid₁ pid₂ m j =>
(m (⟨1, by omega⟩ : Fin 2) + (pid₂ * BLK + j.val / DM)) * rs + pid₁ * hs + j.val % DM
mask1 := fun _ _ pid₂ m _ j => pid₂ * BLK + j.val / DM < m (⟨0, by omega⟩ : Fin 2)
mask2 := fun _ _ _ m t j => t.val * BLK + j.val % BLK < m (⟨0, by omega⟩ : Fin 2)
mask3 := fun _ _ _ m t j => t.val * BLK + j.val / DM < m (⟨0, by omega⟩ : Fin 2)
writeMask := fun _ _ pid₂ m j => pid₂ * BLK + j.val / DM < m (⟨0, by omega⟩ : Fin 2)
Referenced model/specification definitions: Python source |
|---|
context_attn_nopadcontext_attn_nopad_output_summary_generalMathematical execution | Configured model / stage Variable-length context attention with equal query/key blocks and shared contiguous layouts; metadata and stream-size assumptions are explicit. | Statement and preconditionsspecification context_attn_nopad_output_summary_general
(Q K V : RegionName) (B_Start_Loc B_Seqlen : Region .nat)
(Out : RegionName) (sm_scale : ℝ) (rs hs BLK DM : Nat)
(hBLK : 0 < BLK) (hDM : 0 < DM) (hDMrs : DM ≤ rs)
(s : BlockState) (hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := context_attn_nopad_fwd_kernel_surface Q K V sm_scale B_Start_Loc B_Seqlen Out
rs hs 1 rs hs 1 rs hs 1 rs hs 1 BLK DM BLK)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLK, DM] => activeG s B_Seqlen BLK idx)
(fun idx : TileIndex [BLK, DM] =>
(Out, outOffsetG s B_Start_Loc rs hs BLK DM idx)))
(expected := fun idx : TileIndex [BLK, DM] =>
ctxNopadGenuineOutValueG s Q K V B_Start_Loc B_Seqlen sm_scale rs hs BLK DM idx)
Referenced model/specification definitions: Python source |
|---|
context_attn_nopadcontext_attn_nopad_io_correctnessAbstract cast/store rounding | Configured model / stage Variable-length context attention with equal query/key blocks and shared contiguous layouts; metadata and stream-size assumptions are explicit. | Statement and preconditionsspecification context_attn_nopad_io_correctness (R : RoundingModel)
(Q K V : RegionName) (B_Start_Loc B_Seqlen : Region .nat) (Out : RegionName)
(sm_scale : ℝ) (rs hs BLK DM NT : Nat)
(hBLK : 0 < BLK) (hDM : 0 < DM) (hDMrs : DM ≤ rs) (hNT : 0 < NT) :
contextAttnNopadIO Q K V B_Start_Loc B_Seqlen Out sm_scale rs hs BLK DM NT ⊨[R]
fun _ _ pid₂ m xs ys zs j =>
contextAttnNopadIOSpec BLK DM NT (m (⟨0, by omega⟩ : Fin 2)) hBLK hNT sm_scale
pid₂ xs ys zs j
IO definition: contextAttnNopadIOdef contextAttnNopadIO (Q K V : RegionName) (B_Start_Loc B_Seqlen : Region .nat)
(Out : RegionName) (sm_scale : ℝ) (rs hs BLK DM NT : Nat) :
StreamMetaMasked3DKernelIO₃ where
kernel := context_attn_nopad_fwd_kernel_surface Q K V sm_scale B_Start_Loc B_Seqlen Out
rs hs 1 rs hs 1 rs hs 1 rs hs 1 BLK DM BLK
inp1 := Q
inp2 := K
inp3 := V
out := Out
nMeta := 2
sty := fun _ => ChanTy.nat
mbuf := ctxNopadMetaBuf B_Start_Loc B_Seqlen
mwin := fun _ pid₀ _ _ => pid₀
T := NT
B1 := BLK * DM
B2 := DM * BLK
B3 := BLK * DM
C := BLK * DM
pre := fun _ _ pid₂ m => pid₂ < NT ∧ m (⟨0, by omega⟩ : Fin 2) ≤ NT * BLK
read1 := fun _ pid₁ pid₂ m _ j =>
(m (⟨1, by omega⟩ : Fin 2) + (pid₂ * BLK + j.val / DM)) * rs + pid₁ * hs + j.val % DM
read2 := fun _ pid₁ _ m t j =>
(m (⟨1, by omega⟩ : Fin 2) + (t.val * BLK + j.val % BLK)) * rs + pid₁ * hs + j.val / BLK
read3 := fun _ pid₁ _ m t j =>
(m (⟨1, by omega⟩ : Fin 2) + (t.val * BLK + j.val / DM)) * rs + pid₁ * hs + j.val % DM
write := fun _ pid₁ pid₂ m j =>
(m (⟨1, by omega⟩ : Fin 2) + (pid₂ * BLK + j.val / DM)) * rs + pid₁ * hs + j.val % DM
mask1 := fun _ _ pid₂ m _ j => pid₂ * BLK + j.val / DM < m (⟨0, by omega⟩ : Fin 2)
mask2 := fun _ _ _ m t j => t.val * BLK + j.val % BLK < m (⟨0, by omega⟩ : Fin 2)
mask3 := fun _ _ _ m t j => t.val * BLK + j.val / DM < m (⟨0, by omega⟩ : Fin 2)
writeMask := fun _ _ pid₂ m j => pid₂ * BLK + j.val / DM < m (⟨0, by omega⟩ : Fin 2)
Referenced model/specification definitions: Python source |
|---|
cosine_computecos_func_correctnessMathematical execution | Original-kernel candidate Masked per-program cosine over mathematical real values. The theorem does not characterize a hardware cosine approximation or host launch. | Statement and preconditionsspecification cos_func_correctness
(a b : RegionName) (n_elements BLOCK_SIZE : Nat) :
cosIO a b n_elements BLOCK_SIZE ⊨ fun xs i => Real.cos (xs i)
IO definition: cosIOdef cosIO (a b : RegionName) (n_elements BLOCK_SIZE : Nat) :
MaskedKernelIO₁ where
kernel := cos_func a b n_elements BLOCK_SIZE
inp := a
out := b
B := BLOCK_SIZE
read := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < n_elements
Referenced model/specification definitions: Python source |
|---|
cross_entropy1cross_entropy_fwd_correctnessMathematical execution | Configured model / stage Forward loss/LSE for the stated smoothing and split modes, with per-block outputs. The Python backward kernel and host reduction of split outputs are outside this headline. | Statement and preconditionsspecification cross_entropy_fwd_correctness
(loss_ptr lse_ptr logits_ptr : RegionName) (labels_ptr : Region .int)
(smoothing lse_square_scale : ℝ) (ignored_index : Int)
(total_classes : Nat) (class_start_idx : Int)
(n_cols n_rows logits_row_stride BLOCK_SIZE : Nat)
(HAS_SMOOTHING SPLIT : Bool)
(hB : 0 < BLOCK_SIZE)
(hne : lse_ptr ≠ loss_ptr) (hLL : lse_ptr ≠ logits_ptr) :
crossEntropyFwdIO loss_ptr lse_ptr logits_ptr labels_ptr smoothing
lse_square_scale ignored_index total_classes class_start_idx n_cols
n_rows logits_row_stride BLOCK_SIZE HAS_SMOOTHING SPLIT ⊨
fun _ pid₁ lab xs g =>
if pid₁ * BLOCK_SIZE < n_cols then
(ceLossLocal n_cols total_classes BLOCK_SIZE smoothing
lse_square_scale ignored_index class_start_idx HAS_SMOOTHING SPLIT
pid₁ lab xs g,
ceBlockLSE n_cols BLOCK_SIZE pid₁ xs)
else (0, 0)
IO definition: crossEntropyFwdIOdef crossEntropyFwdIO
(loss_ptr lse_ptr logits_ptr : RegionName) (labels_ptr : Region .int)
(smoothing lse_square_scale : ℝ) (ignored_index : Int)
(total_classes : Nat) (class_start_idx : Int)
(n_cols n_rows logits_row_stride BLOCK_SIZE : Nat)
(HAS_SMOOTHING SPLIT : Bool) : MetaGatherMasked2DKernelIO₂ₓ₂ where
kernel := cross_entropy_fwd_surface loss_ptr lse_ptr logits_ptr labels_ptr
smoothing lse_square_scale ignored_index total_classes class_start_idx
n_cols n_rows logits_row_stride BLOCK_SIZE HAS_SMOOTHING SPLIT
mbufL := Region.cast labels_ptr
inp := logits_ptr
out1 := loss_ptr
out2 := lse_ptr
B := BLOCK_SIZE
mwinL := fun pid₀ _ => pid₀
read := fun pid₀ pid₁ _ j =>
pid₀ * logits_row_stride + (pid₁ * BLOCK_SIZE + j.val)
mask := fun _ pid₁ _ j => pid₁ * BLOCK_SIZE + j.val < n_cols
gwin := fun pid₀ _ lab =>
pid₀ * logits_row_stride + (lab - class_start_idx).toNat
gmask := fun _ pid₁ lab =>
lab ≠ ignored_index ∧
lab - class_start_idx ≥ (↑(pid₁ * BLOCK_SIZE) : Int) ∧
lab - class_start_idx < (↑(min n_cols ((pid₁ + 1) * BLOCK_SIZE)) : Int)
write1 := fun pid₀ pid₁ _ => pid₁ * n_rows + pid₀
write2 := fun pid₀ pid₁ _ => pid₁ * n_rows + pid₀
Referenced model/specification definitions: Python source |
|---|
cross_entropy2cross_entropy_fwd_correctnessMathematical execution | Configured model / stage Forward loss/LSE/z-loss for explicit smoothing, scaling, and split modes. The backward kernel and host aggregation are not established by this forward theorem. | Statement and preconditionsspecification cross_entropy_fwd_correctness
(loss_ptr lse_ptr z_loss_ptr logits_ptr : RegionName) (labels_ptr : Region .int)
(smoothing logit_scale lse_square_scale : ℝ) (ignored_index : Int)
(total_classes : Nat) (class_start_idx : Int)
(n_cols n_rows logits_row_stride BLOCK_SIZE : Nat)
(HAS_SMOOTHING SPLIT : Bool)
(hB : 0 < BLOCK_SIZE)
(hne : lse_ptr ≠ loss_ptr) (hneZ : lse_ptr ≠ z_loss_ptr)
(hLL : lse_ptr ≠ logits_ptr) (hLZ : loss_ptr ≠ z_loss_ptr) :
crossEntropyFwdIO loss_ptr lse_ptr z_loss_ptr logits_ptr labels_ptr smoothing
logit_scale lse_square_scale ignored_index total_classes class_start_idx
n_cols n_rows logits_row_stride BLOCK_SIZE HAS_SMOOTHING SPLIT ⊨
fun _ pid₁ lab xs g =>
if pid₁ * BLOCK_SIZE < n_cols then
(ceLossLocal n_cols total_classes BLOCK_SIZE smoothing logit_scale
lse_square_scale ignored_index class_start_idx HAS_SMOOTHING SPLIT
pid₁ lab xs g,
ceBlockLSE n_cols BLOCK_SIZE pid₁ logit_scale xs,
ceZLossLocal n_cols BLOCK_SIZE logit_scale lse_square_scale
ignored_index pid₁ lab xs)
else (0, 0, 0)
IO definition: crossEntropyFwdIOdef crossEntropyFwdIO
(loss_ptr lse_ptr z_loss_ptr logits_ptr : RegionName) (labels_ptr : Region .int)
(smoothing logit_scale lse_square_scale : ℝ) (ignored_index : Int)
(total_classes : Nat) (class_start_idx : Int)
(n_cols n_rows logits_row_stride BLOCK_SIZE : Nat)
(HAS_SMOOTHING SPLIT : Bool) : MetaMasked2DKernelIO₂ₓ₃ where
kernel := cross_entropy_fwd_surface loss_ptr lse_ptr z_loss_ptr logits_ptr
labels_ptr smoothing logit_scale lse_square_scale ignored_index total_classes
class_start_idx n_cols n_rows logits_row_stride BLOCK_SIZE HAS_SMOOTHING SPLIT
mbufL := Region.cast labels_ptr
inp := logits_ptr
out1 := loss_ptr
out2 := lse_ptr
out3 := z_loss_ptr
B := BLOCK_SIZE
mwinL := fun pid₀ _ => pid₀
read := fun pid₀ pid₁ _ j =>
pid₀ * logits_row_stride + (pid₁ * BLOCK_SIZE + j.val)
mask := fun _ pid₁ _ j => pid₁ * BLOCK_SIZE + j.val < n_cols
gwin := fun pid₀ _ lab =>
pid₀ * logits_row_stride + (lab - class_start_idx).toNat
gmask := fun _ pid₁ lab =>
lab ≠ ignored_index ∧
lab - class_start_idx ≥ (↑(pid₁ * BLOCK_SIZE) : Int) ∧
lab - class_start_idx < (↑(min n_cols ((pid₁ + 1) * BLOCK_SIZE)) : Int)
write1 := fun pid₀ pid₁ _ => pid₁ * n_rows + pid₀
write2 := fun pid₀ pid₁ _ => pid₁ * n_rows + pid₀
write3 := fun pid₀ pid₁ _ => pid₁ * n_rows + pid₀
writeMask3 := fun _ _ _ => SPLIT = Bool.false
Referenced model/specification definitions: Python source |
|---|
cross_entropy_opscross_entropy_fwd_correctnessMathematical execution | Configured model / stage Forward loss/LSE/z-loss under the listed flags and label/layout premises; separate Python backward execution and host split aggregation are outside the headline. | Statement and preconditionsspecification cross_entropy_fwd_correctness
(loss_ptr lse_ptr z_loss_ptr logits_ptr : RegionName) (labels_ptr : Region .int)
(smoothing logit_scale lse_square_scale : ℝ) (ignored_index : Int)
(total_classes : Nat) (class_start_idx : Int)
(n_cols n_rows logits_row_stride BLOCK_SIZE : Nat)
(HAS_SMOOTHING SPLIT : Bool)
(hB : 0 < BLOCK_SIZE)
(hne : lse_ptr ≠ loss_ptr) (hneZ : lse_ptr ≠ z_loss_ptr)
(hLL : lse_ptr ≠ logits_ptr) (hLZ : loss_ptr ≠ z_loss_ptr) :
crossEntropyFwdIO loss_ptr lse_ptr z_loss_ptr logits_ptr labels_ptr smoothing
logit_scale lse_square_scale ignored_index total_classes class_start_idx
n_cols n_rows logits_row_stride BLOCK_SIZE HAS_SMOOTHING SPLIT ⊨
fun _ pid₁ lab xs g =>
if pid₁ * BLOCK_SIZE < n_cols then
(ceLossLocal n_cols total_classes BLOCK_SIZE smoothing logit_scale
lse_square_scale ignored_index class_start_idx HAS_SMOOTHING SPLIT
pid₁ lab xs g,
ceBlockLSE n_cols BLOCK_SIZE pid₁ logit_scale xs,
ceZLossLocal n_cols BLOCK_SIZE logit_scale lse_square_scale
ignored_index pid₁ lab xs)
else (0, 0, 0)
IO definition: crossEntropyFwdIOdef crossEntropyFwdIO
(loss_ptr lse_ptr z_loss_ptr logits_ptr : RegionName) (labels_ptr : Region .int)
(smoothing logit_scale lse_square_scale : ℝ) (ignored_index : Int)
(total_classes : Nat) (class_start_idx : Int)
(n_cols n_rows logits_row_stride BLOCK_SIZE : Nat)
(HAS_SMOOTHING SPLIT : Bool) : MetaMasked2DKernelIO₂ₓ₃ where
kernel := cross_entropy_fwd_surface loss_ptr lse_ptr z_loss_ptr logits_ptr
labels_ptr smoothing logit_scale lse_square_scale ignored_index total_classes
class_start_idx n_cols n_rows logits_row_stride BLOCK_SIZE HAS_SMOOTHING SPLIT
mbufL := Region.cast labels_ptr
inp := logits_ptr
out1 := loss_ptr
out2 := lse_ptr
out3 := z_loss_ptr
B := BLOCK_SIZE
mwinL := fun pid₀ _ => pid₀
read := fun pid₀ pid₁ _ j =>
pid₀ * logits_row_stride + (pid₁ * BLOCK_SIZE + j.val)
mask := fun _ pid₁ _ j => pid₁ * BLOCK_SIZE + j.val < n_cols
gwin := fun pid₀ _ lab =>
pid₀ * logits_row_stride + (lab - class_start_idx).toNat
gmask := fun _ pid₁ lab =>
lab ≠ ignored_index ∧
lab - class_start_idx ≥ (↑(pid₁ * BLOCK_SIZE) : Int) ∧
lab - class_start_idx < (↑(min n_cols ((pid₁ + 1) * BLOCK_SIZE)) : Int)
write1 := fun pid₀ pid₁ _ => pid₁ * n_rows + pid₀
write2 := fun pid₀ pid₁ _ => pid₁ * n_rows + pid₀
write3 := fun pid₀ pid₁ _ => pid₁ * n_rows + pid₀
writeMask3 := fun _ _ _ => SPLIT = Bool.false
Referenced model/specification definitions: Python source |
|---|
decay_cumsumdecay_cumsum_backward_closed_output_summary_generalMathematical execution | Configured model / stage Separate forward-decay, QG/KG preparation, and backward-DG contracts. Stride, tile-size, and region-separation premises are explicit; this does not compose a whole host pipeline. | Statement and preconditionsspecification decay_cumsum_backward_closed_output_summary_general :
(ComputeCorrect.Realizes_without_Rounding
(kernel := bwd_decay_global_cumsum_surface DQInner DQInter DKInner DKInter
Q K G DG s_qk_h DK BT BK)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf (active s DK BK)
(fun i => (DQInter, offset s s_qk_h DK t_rel.val BT BK i)))
(expected := fun i : Fin BK =>
bwdDQInterClosed s DQInner DQInter G s_qk_h DK BT BK t_rel i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := bwd_decay_global_cumsum_surface DQInner DQInter DKInner DKInter
Q K G DG s_qk_h DK BT BK)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf (active s DK BK)
(fun i => (DKInter, offset s s_qk_h DK t_rel.val BT BK i)))
(expected := fun i : Fin BK =>
bwdDKInterClosed s DKInner DKInter G s_qk_h DK BT BK t_rel i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := bwd_decay_global_cumsum_surface DQInner DQInter DKInner DKInter
Q K G DG s_qk_h DK BT BK)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf (active s DK BK)
(fun i => (DG, offset s s_qk_h DK t_rel.val BT BK i)))
(expected := fun i : Fin BK =>
bwdDGClosed s DQInner DQInter DKInner DKInter Q K G s_qk_h DK BT BK t_rel i))
Referenced model/specification definitions: Python source |
|---|
decay_cumsumfwd_decay_cumsum_io_correctnessAbstract cast/store rounding | Configured model / stage Separate forward-decay, QG/KG preparation, and backward-DG contracts. Stride, tile-size, and region-separation premises are explicit; this does not compose a whole host pipeline. | Statement and preconditionsspecification fwd_decay_cumsum_io_correctness (R : RoundingModel)
(G GO : RegionName) (s_qk_h s_qk_t s_qk_d B H T : Nat) (scale : ℝ)
(BT BK DK : Nat) (hne : G ≠ GO) (hBK : BK ≤ DK) :
fwdDecayCumsumKernelIO G GO s_qk_h s_qk_t s_qk_d B H T scale BT BK DK ⊨[R]
fun _ _ _ gs t j => fwdDecayStreamSpec BT BK gs t j
IO definition: fwdDecayCumsumKernelIOdef fwdDecayCumsumKernelIO (G GO : RegionName)
(s_qk_h s_qk_t s_qk_d B H T : Nat) (scale : ℝ) (BT BK DK : Nat) :
StreamEmitMasked3DKernelIO₁ where
kernel := fwd_decay_cumsum_surface G GO s_qk_h s_qk_t s_qk_d B H T scale BT BK DK
inp1 := G
out := GO
T := BT
B1 := BK
C := BK
read1 := fun p₀ p₁ p₂ t j => p₂ * s_qk_h + p₁ * BT * DK + p₀ * BK + j.val + t.val * DK
write := fun p₀ p₁ p₂ t j => p₂ * s_qk_h + p₁ * BT * DK + p₀ * BK + j.val + t.val * DK
mask1 := fun p₀ _ _ _ j => p₀ * BK + j.val < DK
writeMask := fun p₀ _ _ _ j => p₀ * BK + j.val < DK
Referenced model/specification definitions: Python source |
|---|
decay_cumsumprepare_qg_kg_io_correctnessAbstract cast/store rounding | Configured model / stage Separate forward-decay, QG/KG preparation, and backward-DG contracts. Stride, tile-size, and region-separation premises are explicit; this does not compose a whole host pipeline. | Statement and preconditionsspecification prepare_qg_kg_io_correctness (R : RoundingModel)
(Q K G QG KG : RegionName) (s_qk_h DK BT BK : Nat) (scale : ℝ)
(hQ_QG : Q ≠ QG) (hQ_KG : Q ≠ KG) (hK_QG : K ≠ QG) (hK_KG : K ≠ KG)
(hG_QG : G ≠ QG) (hG_KG : G ≠ KG) (hQG_KG : QG ≠ KG)
(hBK : BK ≤ DK) (hBT : 0 < BT) :
prepareQgKgKernelIO Q K G QG KG s_qk_h DK BT BK scale ⊨[R]
fun _ _ _ xs o t j => prepStreamSpec BT BK scale xs o t j
IO definition: prepareQgKgKernelIOdef prepareQgKgKernelIO (Q K G QG KG : RegionName)
(s_qk_h DK BT BK : Nat) (scale : ℝ) :
StreamGroupedEmitMasked3DKernelIO where
kernel := prepare_qg_kg_surface Q K G QG KG s_qk_h DK BT BK scale
nIn := 3
nOut := 2
bufs := [Q, K, G, QG, KG]
inp := fun i => match i with
| ⟨0, _⟩ => Q
| ⟨1, _⟩ => K
| ⟨2, _⟩ => G
| ⟨n + 3, h⟩ => absurd h (by omega)
out := fun o => match o with
| ⟨0, _⟩ => QG
| ⟨1, _⟩ => KG
| ⟨n + 2, h⟩ => absurd h (by omega)
T := BT
B := BK
read := fun _ p₀ p₁ p₂ t j =>
p₂ * s_qk_h + p₁ * BT * DK + p₀ * BK + j.val + t.val * DK
readMask := fun i p₀ _ _ t j => match i with
| ⟨0, _⟩ => p₀ * BK + j.val < DK
| ⟨1, _⟩ => p₀ * BK + j.val < DK
| ⟨2, _⟩ => p₀ * BK + j.val < DK ∨ t.val = BT - 1
| ⟨n + 3, h⟩ => absurd h (by omega)
write := fun _ p₀ p₁ p₂ t j =>
p₂ * s_qk_h + p₁ * BT * DK + p₀ * BK + j.val + t.val * DK
writeMask := fun _ p₀ _ _ _ j => p₀ * BK + j.val < DK
Referenced model/specification definitions: Python source |
|---|
decay_cumsumbwd_decay_global_cumsum_io_correctnessAbstract cast/store rounding | Configured model / stage Separate forward-decay, QG/KG preparation, and backward-DG contracts. Stride, tile-size, and region-separation premises are explicit; this does not compose a whole host pipeline. | Statement and preconditionsspecification bwd_decay_global_cumsum_io_correctness (R : RoundingModel)
(DQInner DQInter DKInner DKInter Q K G DG : RegionName)
(s_qk_h DK BT BK : Nat)
(hDKInner_DQInter : DKInner ≠ DQInter) (hDKInter_DQInter : DKInter ≠ DQInter)
(hQ_DQInter : Q ≠ DQInter) (hQ_DKInter : Q ≠ DKInter)
(hK_DQInter : K ≠ DQInter) (hK_DKInter : K ≠ DKInter)
(hDQInter_DKInter : DQInter ≠ DKInter)
(hDQInter_DG : DQInter ≠ DG) (hDKInter_DG : DKInter ≠ DG)
(hG_DQInter : G ≠ DQInter) (hG_DKInter : G ≠ DKInter) (hG_DG : G ≠ DG)
(hQ_DG : Q ≠ DG) (hK_DG : K ≠ DG)
(hDQInner_DQInter : DQInner ≠ DQInter) (hDQInner_DKInter : DQInner ≠ DKInter)
(hDQInner_DG : DQInner ≠ DG)
(hDKInner_DKInter : DKInner ≠ DKInter) (hDKInner_DG : DKInner ≠ DG)
(hBK : BK ≤ DK) (hBT : 0 < BT) :
bwdDecayCumsumKernelIO DQInner DQInter DKInner DKInter Q K G DG
s_qk_h DK BT BK ⊨[R]
fun _ _ _ xs o t j => bwdDecayStreamSpec BT BK xs o t j
IO definition: bwdDecayCumsumKernelIOdef bwdDecayCumsumKernelIO
(DQInner DQInter DKInner DKInter Q K G DG : RegionName)
(s_qk_h DK BT BK : Nat) : StreamGroupedEmitMasked3DKernelIO where
kernel := bwd_decay_global_cumsum_surface DQInner DQInter DKInner DKInter Q K G DG
s_qk_h DK BT BK
nIn := 7
nOut := 3
bufs := [DQInner, DQInter, DKInner, DKInter, Q, K, G, DG]
inp := fun i => match i with
| ⟨0, _⟩ => G
| ⟨1, _⟩ => DQInner
| ⟨2, _⟩ => DQInter
| ⟨3, _⟩ => DKInner
| ⟨4, _⟩ => DKInter
| ⟨5, _⟩ => Q
| ⟨6, _⟩ => K
| ⟨n + 7, h⟩ => absurd h (by omega)
out := fun o => match o with
| ⟨0, _⟩ => DQInter
| ⟨1, _⟩ => DKInter
| ⟨2, _⟩ => DG
| ⟨n + 3, h⟩ => absurd h (by omega)
T := BT
B := BK
read := fun _ p₀ p₁ p₂ t j =>
p₂ * s_qk_h + p₁ * BT * DK + p₀ * BK + j.val + t.val * DK
readMask := fun _ p₀ _ _ _ j => p₀ * BK + j.val < DK
write := fun _ p₀ p₁ p₂ t j =>
p₂ * s_qk_h + p₁ * BT * DK + p₀ * BK + j.val + t.val * DK
writeMask := fun _ p₀ _ _ _ j => p₀ * BK + j.val < DK
Referenced model/specification definitions: Python source |
|---|
dequantize_matmuldequantize_kernel_correctnessMathematical execution | Configured model / stage The dequantize_kernel stage scales the modeled input tile by its loaded scales under injective output addressing; a surrounding matrix multiplication is not part of this theorem. | Statement and preconditionsspecification dequantize_kernel_correctness
(b_ptr b_scale_ptr fpb_ptr : RegionName)
(K N stride_bk stride_bn stride_fpbk stride_fpbn BLOCK_SIZE_N BLOCK_SIZE_K : Nat)
(hBK : 0 < BLOCK_SIZE_K)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun j : Fin (BLOCK_SIZE_K * BLOCK_SIZE_N) =>
fpbOffset p₀ p₁ stride_fpbk stride_fpbn BLOCK_SIZE_N BLOCK_SIZE_K
(laneIdx BLOCK_SIZE_K BLOCK_SIZE_N j))) :
dequantizeIO b_ptr b_scale_ptr fpb_ptr K N stride_bk stride_bn
stride_fpbk stride_fpbn BLOCK_SIZE_N BLOCK_SIZE_K
⊨ fun _ _ xs ys j => xs j * ys j
IO definition: dequantizeIOdef dequantizeIO (b_ptr b_scale_ptr fpb_ptr : RegionName)
(K N stride_bk stride_bn stride_fpbk stride_fpbn BLOCK_SIZE_N BLOCK_SIZE_K : Nat) :
Masked2DKernelIO₂ where
kernel := dequantize_kernel b_ptr b_scale_ptr fpb_ptr K N stride_bk stride_bn
stride_fpbk stride_fpbn BLOCK_SIZE_N BLOCK_SIZE_K
in1 := b_ptr
in2 := b_scale_ptr
out := fpb_ptr
B := BLOCK_SIZE_K * BLOCK_SIZE_N
read1 := fun p₀ p₁ j =>
bOffset p₀ p₁ stride_bk stride_bn BLOCK_SIZE_N BLOCK_SIZE_K
(laneIdx BLOCK_SIZE_K BLOCK_SIZE_N j)
read2 := fun _ p₁ j =>
nOffset p₁ BLOCK_SIZE_N (laneIdx BLOCK_SIZE_K BLOCK_SIZE_N j).2.1
write := fun p₀ p₁ j =>
fpbOffset p₀ p₁ stride_fpbk stride_fpbn BLOCK_SIZE_N BLOCK_SIZE_K
(laneIdx BLOCK_SIZE_K BLOCK_SIZE_N j)
mask := fun p₀ p₁ j =>
dequantizeActive p₀ p₁ K N BLOCK_SIZE_N BLOCK_SIZE_K
(laneIdx BLOCK_SIZE_K BLOCK_SIZE_N j)
read2Mask := fun _ p₁ j =>
nOffset p₁ BLOCK_SIZE_N (laneIdx BLOCK_SIZE_K BLOCK_SIZE_N j).2.1 < N
Referenced model/specification definitions: Python source |
|---|
dequantize_rowwisedequantize_rowwise_correctnessMathematical execution | Configured model / stage One masked rowwise dequantization program with positive BLOCK_SIZE/P2 and a supplied inv_127 parameter; mathematical values do not certify integer storage formats. | Statement and preconditionsspecification dequantize_rowwise_correctness
(x_ptr state_x output_ptr : RegionName)
(inv_127 : ℝ) (n_elements BLOCK_SIZE P2 : Nat)
(hBS : 0 < BLOCK_SIZE) (hP2 : 0 < P2) :
dequantizeRowwiseIO x_ptr state_x output_ptr inv_127 n_elements
BLOCK_SIZE P2 ⊨ fun _ _ xs ys i => ys i * xs i * inv_127
IO definition: dequantizeRowwiseIOdef dequantizeRowwiseIO (x_ptr state_x output_ptr : RegionName)
(inv_127 : ℝ) (n_elements BLOCK_SIZE P2 : Nat) : Masked2DKernelIO₂ where
kernel := dequantize_rowwise_kernel x_ptr state_x output_ptr inv_127
n_elements BLOCK_SIZE P2
in1 := x_ptr
in2 := state_x
out := output_ptr
B := P2
read1 := fun pid _ j => pid * BLOCK_SIZE + j.val
read2 := fun pid _ _ => pid
write := fun pid _ j => pid * BLOCK_SIZE + j.val
mask := fun _ _ j => j.val < BLOCK_SIZE
Referenced model/specification definitions: Python source |
|---|
destindex_copydestindex_copy_correctnessMathematical execution | Configured model / stage Two supplied KV streams are copied to destinations selected by natural-valued metadata. Bounds and address-disjointness premises exclude invalid or colliding writes. | Statement and preconditionsspecification destindex_copy_correctness
(KV_nope KV_rope : RegionName) (Dest_loc : Region .nat) (O_nope O_rope : RegionName)
(stride_kv_nope_bs stride_kv_nope_h stride_kv_nope_d
stride_kv_rope_bs stride_kv_rope_h stride_kv_rope_d
stride_o_nope_bs stride_o_nope_h stride_o_nope_d
stride_o_rope_bs stride_o_rope_h stride_o_rope_d
kv_nope_head_num kv_rope_head_num
BLOCK_DMODEL_NOPE BLOCK_DMODEL_ROPE : Nat)
(hRegion : O_nope ≠ O_rope)
(hNopeInj : Function.Injective
(fun j : Fin BLOCK_DMODEL_NOPE => stride_o_nope_d * j.val))
(hRopeInj : Function.Injective
(fun j : Fin BLOCK_DMODEL_ROPE => stride_o_rope_d * j.val)) :
destindexCopyIO KV_nope KV_rope Dest_loc O_nope O_rope
stride_kv_nope_bs stride_kv_nope_h stride_kv_nope_d
stride_kv_rope_bs stride_kv_rope_h stride_kv_rope_d
stride_o_nope_bs stride_o_nope_h stride_o_nope_d
stride_o_rope_bs stride_o_rope_h stride_o_rope_d
kv_nope_head_num kv_rope_head_num
BLOCK_DMODEL_NOPE BLOCK_DMODEL_ROPE
⊨ fun _ _ _ xs ys => (xs, ys)
IO definition: destindexCopyIOdef destindexCopyIO
(KV_nope KV_rope : RegionName) (Dest_loc : Region .nat) (O_nope O_rope : RegionName)
(stride_kv_nope_bs stride_kv_nope_h stride_kv_nope_d
stride_kv_rope_bs stride_kv_rope_h stride_kv_rope_d
stride_o_nope_bs stride_o_nope_h stride_o_nope_d
stride_o_rope_bs stride_o_rope_h stride_o_rope_d
kv_nope_head_num kv_rope_head_num
BLOCK_DMODEL_NOPE BLOCK_DMODEL_ROPE : Nat) :
MetaMasked2DKernelIO₂ₓ₂ where
kernel := fwd_kernel_destindex_copy_kv KV_nope KV_rope Dest_loc O_nope O_rope
stride_kv_nope_bs stride_kv_nope_h stride_kv_nope_d
stride_kv_rope_bs stride_kv_rope_h stride_kv_rope_d
stride_o_nope_bs stride_o_nope_h stride_o_nope_d
stride_o_rope_bs stride_o_rope_h stride_o_rope_d
kv_nope_head_num kv_rope_head_num
BLOCK_DMODEL_NOPE BLOCK_DMODEL_ROPE
mbuf1 := Dest_loc
in1 := KV_nope
in2 := KV_rope
out1 := O_nope
out2 := O_rope
B1 := BLOCK_DMODEL_NOPE
B2 := BLOCK_DMODEL_ROPE
mwin1 := fun pid₀ _ => pid₀
read1 := fun pid₀ _ _ j => pid₀ * stride_kv_nope_bs + stride_kv_nope_d * j.val
read2 := fun pid₀ _ _ j => pid₀ * stride_kv_rope_bs + stride_kv_rope_d * j.val
write1 := fun _ _ m₁ j => m₁ * stride_o_nope_bs + stride_o_nope_d * j.val
write2 := fun _ _ m₁ j => m₁ * stride_o_rope_bs + stride_o_rope_d * j.val
Referenced model/specification definitions: Python source |
|---|
destindex_copy_kv1fwd_kernel_destindex_copy_kv_correctnessMathematical execution | Configured model / stage Masked destination-index KV copy with explicit head dimension, strides, and injective destination addresses; host metadata generation is external. | Statement and preconditionsspecification fwd_kernel_destindex_copy_kv_correctness
(K : RegionName) (Dest_loc : Region .nat) (Out : RegionName)
(stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
head_num head_dim BLOCK_DMODEL BLOCK_HEAD : Nat)
(hInj : Function.Injective
(fun idx : TileIndex [BLOCK_HEAD, BLOCK_DMODEL] =>
stride_o_h * headIndex idx + stride_o_d * dimIndex idx)) :
destindexCopyKvIO K Dest_loc Out
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
head_num head_dim BLOCK_DMODEL BLOCK_HEAD
⊨ fun _ _ _ _ xs => xs
IO definition: destindexCopyKvIOdef destindexCopyKvIO
(K : RegionName) (Dest_loc : Region .nat) (Out : RegionName)
(stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
head_num head_dim BLOCK_DMODEL BLOCK_HEAD : Nat) :
MetaMasked2DKernelIO₁ where
kernel := fwd_kernel_destindex_copy_kv K Dest_loc Out
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
head_num head_dim BLOCK_DMODEL BLOCK_HEAD
mbuf1 := Dest_loc
mbuf2 := Dest_loc
inp := K
out := Out
B := BLOCK_HEAD * BLOCK_DMODEL
mwin1 := fun pid₀ _ => pid₀
mwin2 := fun pid₀ _ => pid₀
read := fun pid₀ _ _ _ j =>
pid₀ * stride_k_bs + stride_k_h * (j.val / BLOCK_DMODEL) +
stride_k_d * (j.val % BLOCK_DMODEL)
write := fun _ _ m₁ _ j =>
m₁ * stride_o_bs + stride_o_h * (j.val / BLOCK_DMODEL) +
stride_o_d * (j.val % BLOCK_DMODEL)
mask := fun _ _ _ _ j =>
j.val / BLOCK_DMODEL < head_num ∧ j.val % BLOCK_DMODEL < head_dim
Referenced model/specification definitions: Python source |
|---|
destindex_copy_kv2fwd_kernel_destindex_copy_kv_correctnessMathematical execution | Configured model / stage Destination-index KV copy for the declared block/head layout and injective destination addresses; host metadata validity and launch composition are external. | Statement and preconditionsspecification fwd_kernel_destindex_copy_kv_correctness
(K : RegionName) (Dest_loc : Region .nat) (Out : RegionName)
(stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
head_num BLOCK_DMODEL BLOCK_HEAD : Nat)
(hInj : Function.Injective
(fun idx : TileIndex [BLOCK_HEAD, BLOCK_DMODEL] =>
stride_o_h * headIndex idx + stride_o_d * dimIndex idx)) :
destindexCopyKvIO K Dest_loc Out
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
head_num BLOCK_DMODEL BLOCK_HEAD
⊨ fun _ _ _ _ xs => xs
IO definition: destindexCopyKvIOdef destindexCopyKvIO
(K : RegionName) (Dest_loc : Region .nat) (Out : RegionName)
(stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
head_num BLOCK_DMODEL BLOCK_HEAD : Nat) :
MetaMasked2DKernelIO₁ where
kernel := fwd_kernel_destindex_copy_kv K Dest_loc Out
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
head_num BLOCK_DMODEL BLOCK_HEAD
mbuf1 := Dest_loc
mbuf2 := Dest_loc
inp := K
out := Out
B := BLOCK_HEAD * BLOCK_DMODEL
mwin1 := fun pid₀ _ => pid₀
mwin2 := fun pid₀ _ => pid₀
read := fun pid₀ _ _ _ j =>
pid₀ * stride_k_bs + stride_k_h * (j.val / BLOCK_DMODEL) +
stride_k_d * (j.val % BLOCK_DMODEL)
write := fun _ _ m₁ _ j =>
m₁ * stride_o_bs + stride_o_h * (j.val / BLOCK_DMODEL) +
stride_o_d * (j.val % BLOCK_DMODEL)
mask := fun _ _ _ _ j => j.val / BLOCK_DMODEL < head_num
Referenced model/specification definitions: Python source |
|---|
diag_ssm_tritondiag_ssm_forward_kernel_output_summaryMathematical execution | Configured model / stage Real-valued diagonal SSM forward/backward recurrence contracts with the displayed sequence/block and region-separation premises. The complex-valued Python variants are not covered. | Statement and preconditionsspecification diag_ssm_forward_kernel_output_summary
(s_ptr x_ptr lambda_ptr y_ptr : RegionName)
(length batch_size dim BLOCK_SIZE : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun idx : TileIndex [length, BLOCK_SIZE] =>
diagSsmForwardOutOffset s batch_size dim BLOCK_SIZE idx))
(hXOutNe : x_ptr ≠ y_ptr) :
(∃ alg, (diag_ssm_forward_kernel s_ptr x_ptr lambda_ptr y_ptr
length batch_size dim BLOCK_SIZE).toAlgorithm? = Except.ok alg) ∧
diag_ssm_forward_kernel_correct_target s_ptr x_ptr lambda_ptr y_ptr
length batch_size dim BLOCK_SIZE s
Referenced model/specification definitions: Python source |
|---|
diag_ssm_tritondiag_ssm_forward_io_correctnessAbstract cast/store rounding | Configured model / stage Real-valued diagonal SSM forward/backward recurrence contracts with the displayed sequence/block and region-separation premises. The complex-valued Python variants are not covered. | Statement and preconditionsspecification diag_ssm_forward_io_correctness (R : RoundingModel)
(s_ptr x_ptr lambda_ptr y_ptr : RegionName)
(length batch_size dim BLOCK_SIZE : Nat)
(hT : 0 < length) (hBS : BLOCK_SIZE ≤ batch_size * dim)
(hXOutNe : x_ptr ≠ y_ptr) :
diagSsmForwardKernelIO s_ptr x_ptr lambda_ptr y_ptr length batch_size dim
BLOCK_SIZE ⊨[R]
fun _ _ ss ls xs t j =>
diagSsmStreamState (ss t j) (ls t j) (fun u => xs u j) (t.val + 1)
IO definition: diagSsmForwardKernelIOdef diagSsmForwardKernelIO (s_ptr x_ptr lambda_ptr y_ptr : RegionName)
(length batch_size dim BLOCK_SIZE : Nat) :
StreamEmitMasked2DKernelIO₃ where
kernel := diag_ssm_forward_kernel s_ptr x_ptr lambda_ptr y_ptr length
batch_size dim BLOCK_SIZE
inp1 := s_ptr
inp2 := lambda_ptr
inp3 := x_ptr
out := y_ptr
T := length
B1 := BLOCK_SIZE
B2 := BLOCK_SIZE
B3 := BLOCK_SIZE
C := BLOCK_SIZE
read1 := fun p₀ _ _ j => p₀ * BLOCK_SIZE + j.val
read2 := fun p₀ _ _ j => (p₀ * BLOCK_SIZE + j.val) % dim
read3 := fun p₀ _ t j => t.val * (batch_size * dim) + (p₀ * BLOCK_SIZE + j.val)
write := fun p₀ _ t j => t.val * (batch_size * dim) + (p₀ * BLOCK_SIZE + j.val)
mask1 := fun p₀ _ _ j => p₀ * BLOCK_SIZE + j.val < batch_size * dim
mask2 := fun p₀ _ _ j => p₀ * BLOCK_SIZE + j.val < batch_size * dim
mask3 := fun p₀ _ _ j => p₀ * BLOCK_SIZE + j.val < batch_size * dim
writeMask := fun p₀ _ _ j => p₀ * BLOCK_SIZE + j.val < batch_size * dim
Referenced model/specification definitions: Python source |
|---|
diag_ssm_tritondiag_ssm_backward_io_correctnessAbstract cast/store rounding | Configured model / stage Real-valued diagonal SSM forward/backward recurrence contracts with the displayed sequence/block and region-separation premises. The complex-valued Python variants are not covered. | Statement and preconditionsspecification diag_ssm_backward_io_correctness (R : RoundingModel)
(s_ptr lambda_ptr y_ptr grad_s_ptr grad_x_ptr grad_lambda_ptr
grad_y_ptr : RegionName)
(length batch_size dim BLOCK_SIZE : Nat)
(hT : 0 < length) (hBS : BLOCK_SIZE ≤ batch_size * dim)
(hGradXSNe : grad_x_ptr ≠ s_ptr) (hGradXYNe : grad_x_ptr ≠ y_ptr)
(hGradXGradYNe : grad_x_ptr ≠ grad_y_ptr)
(hGradXLambdaNe : grad_x_ptr ≠ lambda_ptr)
(hGradSGradXNe : grad_s_ptr ≠ grad_x_ptr)
(hGradLambdaGradXNe : grad_lambda_ptr ≠ grad_x_ptr)
(hGradLambdaGradSNe : grad_lambda_ptr ≠ grad_s_ptr) :
diagSsmBackwardKernelIO s_ptr lambda_ptr y_ptr grad_s_ptr grad_x_ptr
grad_lambda_ptr grad_y_ptr length batch_size dim BLOCK_SIZE ⊨[R]
fun _ _ _ xs o t j => diagSsmBackwardStreamSpec xs o t j
IO definition: diagSsmBackwardKernelIOdef diagSsmBackwardKernelIO
(s_ptr lambda_ptr y_ptr grad_s_ptr grad_x_ptr grad_lambda_ptr
grad_y_ptr : RegionName)
(length batch_size dim BLOCK_SIZE : Nat) :
StreamGroupedEmitMasked3DKernelIO where
kernel := diag_ssm_backward_kernel s_ptr lambda_ptr y_ptr grad_s_ptr
grad_x_ptr grad_lambda_ptr grad_y_ptr length batch_size dim BLOCK_SIZE
nIn := 4
nOut := 3
bufs := [s_ptr, lambda_ptr, y_ptr, grad_s_ptr, grad_x_ptr,
grad_lambda_ptr, grad_y_ptr]
inp := fun i => match i with
| ⟨0, _⟩ => grad_y_ptr
| ⟨1, _⟩ => y_ptr
| ⟨2, _⟩ => s_ptr
| ⟨_ + 3, _⟩ => lambda_ptr
out := fun o => match o with
| ⟨0, _⟩ => grad_x_ptr
| ⟨1, _⟩ => grad_s_ptr
| ⟨_ + 2, _⟩ => grad_lambda_ptr
T := length
B := BLOCK_SIZE
read := fun i p₀ _ _ t j => match i with
| ⟨0, _⟩ => t.val * (batch_size * dim) + (p₀ * BLOCK_SIZE + j.val)
| ⟨1, _⟩ => (t.val - 1) * (batch_size * dim) + (p₀ * BLOCK_SIZE + j.val)
| ⟨2, _⟩ => p₀ * BLOCK_SIZE + j.val
| ⟨_ + 3, _⟩ => (p₀ * BLOCK_SIZE + j.val) % dim
readMask := fun i p₀ _ _ t j => match i with
| ⟨0, _⟩ => p₀ * BLOCK_SIZE + j.val < batch_size * dim
| ⟨1, _⟩ => 0 < t.val ∧ p₀ * BLOCK_SIZE + j.val < batch_size * dim
| ⟨2, _⟩ => t.val = 0 ∧ p₀ * BLOCK_SIZE + j.val < batch_size * dim
| ⟨_ + 3, _⟩ => p₀ * BLOCK_SIZE + j.val < batch_size * dim
write := fun o p₀ _ _ t j => match o with
| ⟨0, _⟩ => t.val * (batch_size * dim) + (p₀ * BLOCK_SIZE + j.val)
| ⟨1, _⟩ => p₀ * BLOCK_SIZE + j.val
| ⟨_ + 2, _⟩ => p₀ * BLOCK_SIZE + j.val
writeMask := fun o p₀ _ _ t j => match o with
| ⟨0, _⟩ => p₀ * BLOCK_SIZE + j.val < batch_size * dim
| ⟨1, _⟩ => t.val = 0 ∧ p₀ * BLOCK_SIZE + j.val < batch_size * dim
| ⟨_ + 2, _⟩ => t.val = 0 ∧ p₀ * BLOCK_SIZE + j.val < batch_size * dim
Referenced model/specification definitions: Python source |
|---|
dropout_tritondropout_kernel_correctnessMathematical execution | Configured model / stage Dropout consumes an explicit keep-mask input and applies the mathematical scaling formula. It does not prove random mask generation, RNG distribution, or concrete floating behavior. | Statement and preconditionsspecification dropout_kernel_correctness
(x_ptr x_keep_ptr output_ptr : RegionName)
(n_elements : Nat) (p : ℝ) (BLOCK_SIZE : Nat) :
dropoutIO x_ptr x_keep_ptr output_ptr n_elements p BLOCK_SIZE ⊨
fun _ _ bs xs i => if bs i then xs i / (1 - p) else 0
IO definition: dropoutIOdef dropoutIO (x_ptr x_keep_ptr output_ptr : RegionName)
(n_elements : Nat) (p : ℝ) (BLOCK_SIZE : Nat) : BoolMasked2DKernelIO₁ where
kernel := dropout_kernel x_ptr x_keep_ptr output_ptr n_elements p BLOCK_SIZE
inp := x_ptr
mbuf := x_keep_ptr
out := output_ptr
B := BLOCK_SIZE
read := fun pid _ j => pid * BLOCK_SIZE + j.val
readm := fun pid _ j => pid * BLOCK_SIZE + j.val
write := fun pid _ j => pid * BLOCK_SIZE + j.val
mask := fun pid _ j => pid * BLOCK_SIZE + j.val < n_elements
Referenced model/specification definitions: Python source |
|---|
embedding_triton_kernelembedding_kernel_output_summaryMathematical execution | Configured model / stage Selected embedding lookup contracts under vocabulary and address assumptions. Body contracts concern one loop slice; the whole IO contract explicitly fixes BLOCK_NN=1. | Statement and preconditionsspecification embedding_kernel_output_summary
(weight input_ids out : RegionName)
(vob_start_id vob_end_id stride_weight_seq stride_out_seq n_ctx
hiden_size BLOCK_DMODEL BLOCK_N BLOCK_NN : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun idx : TileIndex [BLOCK_N, BLOCK_DMODEL] =>
outOffsetFull s stride_out_seq BLOCK_N idx))
(hOne : BLOCK_NN = 1)
(hInputOutNe : input_ids ≠ out)
(hWeightOutNe : weight ≠ out) :
(∃ alg, (embedding_kernel weight input_ids out
vob_start_id vob_end_id stride_weight_seq stride_out_seq n_ctx
hiden_size BLOCK_DMODEL BLOCK_N BLOCK_NN).toAlgorithm? = Except.ok alg) ∧
embedding_kernel_correct_target weight input_ids out
vob_start_id vob_end_id stride_weight_seq stride_out_seq n_ctx
hiden_size BLOCK_DMODEL BLOCK_N BLOCK_NN s
Referenced model/specification definitions: Python source |
|---|
embedding_triton_kernelembedding_body_io_correctnessMathematical execution | Configured model / stage Selected embedding lookup contracts under vocabulary and address assumptions. Body contracts concern one loop slice; the whole IO contract explicitly fixes BLOCK_NN=1. | Statement and preconditionsspecification embedding_body_io_correctness
(weight input_ids out : RegionName)
(vob_start_id vob_end_id stride_weight_seq stride_out_seq n_ctx
hiden_size BLOCK_DMODEL BLOCK_N BLOCK_NN start_nn : Nat)
(hOutInj : ∀ pid : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_NN, BLOCK_DMODEL] =>
(pid * BLOCK_N + start_nn + idx.1.val) * stride_out_seq
+ idx.2.1.val)) :
bodyIO weight input_ids out vob_start_id vob_end_id stride_weight_seq
stride_out_seq n_ctx hiden_size BLOCK_DMODEL BLOCK_N BLOCK_NN start_nn
⊨ fun _pid ids xs j =>
if vob_start_id ≤ ids (j.1, PUnit.unit) ∧
ids (j.1, PUnit.unit) < vob_end_id then xs j else 0
IO definition: bodyIOdef bodyIO (weight input_ids out : RegionName)
(vob_start_id vob_end_id stride_weight_seq stride_out_seq n_ctx
hiden_size BLOCK_DMODEL BLOCK_N BLOCK_NN start_nn : Nat) :
GatherTileKernelIO where
kernel := embedding_body_slice weight input_ids out vob_start_id vob_end_id
stride_weight_seq stride_out_seq n_ctx hiden_size BLOCK_DMODEL BLOCK_N
BLOCK_NN start_nn
idxbuf := input_ids
inp := weight
out := out
shapeIdx := [BLOCK_NN]
shape := [BLOCK_NN, BLOCK_DMODEL]
readx := fun pid i => pid * BLOCK_N + start_nn + i.1.val
read := fun _pid ids j =>
(ids (j.1, PUnit.unit) - vob_start_id) * stride_weight_seq + j.2.1.val
write := fun pid j =>
(pid * BLOCK_N + start_nn + j.1.val) * stride_out_seq + j.2.1.val
maskx := fun pid i => pid * BLOCK_N + start_nn + i.1.val < n_ctx
readMask := fun _pid ids j =>
(vob_start_id ≤ ids (j.1, PUnit.unit) ∧
ids (j.1, PUnit.unit) < vob_end_id) ∧ j.2.1.val < hiden_size
writeMask := fun pid j =>
pid * BLOCK_N + start_nn + j.1.val < n_ctx ∧ j.2.1.val < hiden_size
Referenced model/specification definitions: Python source |
|---|
embedding_triton_kernelembedding_body_io_correctnessRAbstract cast/store rounding | Configured model / stage Selected embedding lookup contracts under vocabulary and address assumptions. Body contracts concern one loop slice; the whole IO contract explicitly fixes BLOCK_NN=1. | Statement and preconditionsspecification embedding_body_io_correctnessR (R : RoundingModel)
(weight input_ids out : RegionName)
(vob_start_id vob_end_id stride_weight_seq stride_out_seq n_ctx
hiden_size BLOCK_DMODEL BLOCK_N BLOCK_NN start_nn : Nat)
(hOutInj : ∀ pid : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_NN, BLOCK_DMODEL] =>
(pid * BLOCK_N + start_nn + idx.1.val) * stride_out_seq
+ idx.2.1.val)) :
bodyIO weight input_ids out vob_start_id vob_end_id stride_weight_seq
stride_out_seq n_ctx hiden_size BLOCK_DMODEL BLOCK_N BLOCK_NN start_nn
⊨[R, FloatDType.real] fun _pid ids xs j =>
if vob_start_id ≤ ids (j.1, PUnit.unit) ∧
ids (j.1, PUnit.unit) < vob_end_id then xs j else 0
IO definition: bodyIOdef bodyIO (weight input_ids out : RegionName)
(vob_start_id vob_end_id stride_weight_seq stride_out_seq n_ctx
hiden_size BLOCK_DMODEL BLOCK_N BLOCK_NN start_nn : Nat) :
GatherTileKernelIO where
kernel := embedding_body_slice weight input_ids out vob_start_id vob_end_id
stride_weight_seq stride_out_seq n_ctx hiden_size BLOCK_DMODEL BLOCK_N
BLOCK_NN start_nn
idxbuf := input_ids
inp := weight
out := out
shapeIdx := [BLOCK_NN]
shape := [BLOCK_NN, BLOCK_DMODEL]
readx := fun pid i => pid * BLOCK_N + start_nn + i.1.val
read := fun _pid ids j =>
(ids (j.1, PUnit.unit) - vob_start_id) * stride_weight_seq + j.2.1.val
write := fun pid j =>
(pid * BLOCK_N + start_nn + j.1.val) * stride_out_seq + j.2.1.val
maskx := fun pid i => pid * BLOCK_N + start_nn + i.1.val < n_ctx
readMask := fun _pid ids j =>
(vob_start_id ≤ ids (j.1, PUnit.unit) ∧
ids (j.1, PUnit.unit) < vob_end_id) ∧ j.2.1.val < hiden_size
writeMask := fun pid j =>
pid * BLOCK_N + start_nn + j.1.val < n_ctx ∧ j.2.1.val < hiden_size
Referenced model/specification definitions: Python source |
|---|
embedding_triton_kernelembedding_kernel_whole_io_correctnessRAbstract cast/store rounding | Configured model / stage Selected embedding lookup contracts under vocabulary and address assumptions. Body contracts concern one loop slice; the whole IO contract explicitly fixes BLOCK_NN=1. | Statement and preconditionsspecification embedding_kernel_whole_io_correctnessR (R : RoundingModel)
(weight input_ids out : RegionName)
(vob_start_id vob_end_id stride_weight_seq stride_out_seq n_ctx
hiden_size BLOCK_DMODEL BLOCK_N : Nat)
(hInputOutNe : input_ids ≠ out)
(hWeightOutNe : weight ≠ out)
(hOutInj : ∀ pid : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_N, BLOCK_DMODEL] =>
(pid * BLOCK_N + idx.1.val) * stride_out_seq + idx.2.1.val)) :
wholeIO weight input_ids out vob_start_id vob_end_id stride_weight_seq
stride_out_seq n_ctx hiden_size BLOCK_DMODEL BLOCK_N
⊨[R, FloatDType.real] fun _pid ids xs j =>
if vob_start_id ≤ ids (j.1, PUnit.unit) ∧
ids (j.1, PUnit.unit) < vob_end_id then xs j else 0
IO definition: wholeIOdef wholeIO (weight input_ids out : RegionName)
(vob_start_id vob_end_id stride_weight_seq stride_out_seq n_ctx
hiden_size BLOCK_DMODEL BLOCK_N : Nat) : GatherTileKernelIO where
kernel := embedding_kernel weight input_ids out vob_start_id vob_end_id
stride_weight_seq stride_out_seq n_ctx hiden_size BLOCK_DMODEL BLOCK_N 1
idxbuf := input_ids
inp := weight
out := out
shapeIdx := [BLOCK_N]
shape := [BLOCK_N, BLOCK_DMODEL]
readx := fun pid i => pid * BLOCK_N + i.1.val
read := fun _pid ids j =>
(ids (j.1, PUnit.unit) - vob_start_id) * stride_weight_seq + j.2.1.val
write := fun pid j => (pid * BLOCK_N + j.1.val) * stride_out_seq + j.2.1.val
maskx := fun pid i => pid * BLOCK_N + i.1.val < n_ctx
readMask := fun _pid ids j =>
(vob_start_id ≤ ids (j.1, PUnit.unit) ∧
ids (j.1, PUnit.unit) < vob_end_id) ∧ j.2.1.val < hiden_size
writeMask := fun pid j => pid * BLOCK_N + j.1.val < n_ctx ∧ j.2.1.val < hiden_size
Referenced model/specification definitions: Python source |
|---|
f8_conversion_utilsf16_to_f8_io_correctnessAbstract cast/store rounding | Configured model / stage Typed elementwise conversion under an abstract RoundingModel, separately for f8e5 and fp16 output stores. No concrete fp8 encoding, IEEE special-value, or hardware conversion proof. | Statement and preconditionsspecification f16_to_f8_io_correctness (R : RoundingModel)
(Y X : RegionName) (N BLOCK_SIZE : Nat) :
f16ToF8IO Y X N BLOCK_SIZE ⊨[R, .f8e5] fun xs i => xs i
IO definition: f16ToF8IOdef f16ToF8IO (Y X : RegionName) (N BLOCK_SIZE : Nat) : MaskedKernelIO₁ where
kernel := kernel_f16_to_f8 Y X N BLOCK_SIZE
inp := X
out := Y
B := BLOCK_SIZE
read := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < N
Referenced model/specification definitions: Python source |
|---|
f8_conversion_utilsf8_to_f16_io_correctnessAbstract cast/store rounding | Configured model / stage Typed elementwise conversion under an abstract RoundingModel, separately for f8e5 and fp16 output stores. No concrete fp8 encoding, IEEE special-value, or hardware conversion proof. | Statement and preconditionsspecification f8_to_f16_io_correctness (R : RoundingModel)
(Y X : RegionName) (N BLOCK_SIZE : Nat) :
f8ToF16IO Y X N BLOCK_SIZE ⊨[R, .fp16] fun xs i => xs i
IO definition: f8ToF16IOdef f8ToF16IO (Y X : RegionName) (N BLOCK_SIZE : Nat) : MaskedKernelIO₁ where
kernel := kernel_f8_to_f16 Y X N BLOCK_SIZE
inp := X
out := Y
B := BLOCK_SIZE
read := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < N
Referenced model/specification definitions: Python source |
|---|
fast_ce_losscross_entropy_forward_correctnessMathematical execution | Configured model / stage Forward and chunked-forward specializations explicitly disable softcapping. Chunked output is a partial loss/logsumexp stage; the backward kernel and host loss aggregation are external. | Statement and preconditionsspecification cross_entropy_forward_correctness
(logits_ptr loss_ptr logsumexp_ptr : RegionName) (labels_ptr : Region .int)
(VOCAB_SIZE logits_row_stride BLOCK_SIZE : Nat)
(SOFTCAP LOGIT_SCALE : ℝ) (DO_LOGIT_SCALING : Bool)
(hV : 0 < VOCAB_SIZE) (hB : 0 < BLOCK_SIZE)
(hne : logsumexp_ptr ≠ loss_ptr) :
fastCeForwardIO logits_ptr loss_ptr logsumexp_ptr labels_ptr VOCAB_SIZE
logits_row_stride BLOCK_SIZE SOFTCAP LOGIT_SCALE DO_LOGIT_SCALING ⊨
fun _ _ lab xs g =>
(fceLseLocal VOCAB_SIZE BLOCK_SIZE LOGIT_SCALE DO_LOGIT_SCALING xs,
fceLossLocal VOCAB_SIZE BLOCK_SIZE LOGIT_SCALE DO_LOGIT_SCALING
lab xs g)
IO definition: fastCeForwardIOdef fastCeForwardIO
(logits_ptr loss_ptr logsumexp_ptr : RegionName) (labels_ptr : Region .int)
(VOCAB_SIZE logits_row_stride BLOCK_SIZE : Nat)
(SOFTCAP LOGIT_SCALE : ℝ) (DO_LOGIT_SCALING : Bool) :
MetaGatherMasked2DKernelIO₂ₓ₂ where
kernel := cross_entropy_forward_surface logits_ptr loss_ptr logsumexp_ptr
labels_ptr VOCAB_SIZE logits_row_stride BLOCK_SIZE SOFTCAP LOGIT_SCALE
Bool.false DO_LOGIT_SCALING
mbufL := Region.cast labels_ptr
inp := logits_ptr
out1 := logsumexp_ptr
out2 := loss_ptr
B := BLOCK_SIZE
mwinL := fun pid₀ _ => pid₀
read := fun pid₀ _ _ j => pid₀ * logits_row_stride + j.val
mask := fun _ _ _ j => j.val < VOCAB_SIZE
gwin := fun pid₀ _ lab => pid₀ * logits_row_stride + lab.toNat
gmask := fun _ _ lab => lab ≠ -100
write1 := fun pid₀ _ _ => pid₀
write2 := fun pid₀ _ _ => pid₀
Referenced model/specification definitions: Python source |
|---|
fast_ce_losschunked_cross_entropy_forward_output_summaryMathematical execution | Configured model / stage Forward and chunked-forward specializations explicitly disable softcapping. Chunked output is a partial loss/logsumexp stage; the backward kernel and host loss aggregation are external. | Statement and preconditionsspecification chunked_cross_entropy_forward_output_summary
(logits_ptr loss_ptr logsumexp_ptr : RegionName) (labels_ptr : Region .int)
(VOCAB_SIZE N_CHUNKS logits_row_stride : Nat)
(SOFTCAP LOGIT_SCALE : ℝ) (DO_LOGIT_SCALING : Bool)
(n : Nat)
(s : BlockState)
(hchunk : s.pids 1 = 0)
(h_tail : 0 * (n+1) < VOCAB_SIZE)
(hne : logsumexp_ptr ≠ loss_ptr) :
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunked_cross_entropy_forward_surface logits_ptr loss_ptr logsumexp_ptr
labels_ptr VOCAB_SIZE N_CHUNKS logits_row_stride (n+1) SOFTCAP LOGIT_SCALE
Bool.false DO_LOGIT_SCALING)
(initialState := s)
(write := fun _ : PUnit => some (logsumexp_ptr, fceChunkLseOffset s N_CHUNKS))
(expected := fun _ =>
fastCeLseSpec (fastCeRowLogits s logits_ptr logits_row_stride VOCAB_SIZE)
0 h_tail (fun x => if DO_LOGIT_SCALING then LOGIT_SCALE * x else x))) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := chunked_cross_entropy_forward_surface logits_ptr loss_ptr logsumexp_ptr
labels_ptr VOCAB_SIZE N_CHUNKS logits_row_stride (n+1) SOFTCAP LOGIT_SCALE
Bool.false DO_LOGIT_SCALING)
(initialState := s)
(write := fun _ : PUnit => some (loss_ptr, fceOutOffset s))
(expected := fun _ =>
(-1 : ℝ) * fceLabelLogit s logits_ptr labels_ptr logits_row_stride
SOFTCAP LOGIT_SCALE Bool.false DO_LOGIT_SCALING))
Referenced model/specification definitions: Python source |
|---|
fast_layernormlayernorm_forward_correctnessMathematical execution | Configured model / stage Forward row normalization with Y/mean/inverse-variance outputs and the stated bounds/region premises. The Python backward kernel is not proved by this forward headline. | Statement and preconditionsspecification layernorm_forward_correctness
(Y X W bias r mu : RegionName)
(Y_row_stride X_row_stride n_cols BLOCK_SIZE : Nat) (eps : ℝ)
(hB : 0 < BLOCK_SIZE)
(hYr : Y ≠ r) (hYmu : Y ≠ mu) (hRmu : r ≠ mu) :
layernormForwardIO Y X W bias r mu Y_row_stride X_row_stride n_cols eps
BLOCK_SIZE ⊨
fun _ _ xs ws bs =>
(fun i => layernormYSpec n_cols BLOCK_SIZE eps xs ws bs i,
fun _ => invVarFullSpec n_cols BLOCK_SIZE eps xs,
fun _ => meanFullSpec n_cols BLOCK_SIZE xs)
IO definition: layernormForwardIOdef layernormForwardIO (Y X W bias r mu : RegionName)
(Y_row_stride X_row_stride n_cols : Nat)
(eps : ℝ) (BLOCK_SIZE : Nat) : Masked2DKernelIO₃ₓ₃ where
kernel := layernorm_forward Y X W bias r mu Y_row_stride X_row_stride
n_cols eps BLOCK_SIZE
in1 := X
in2 := W
in3 := bias
out1 := Y
out2 := r
out3 := mu
B := BLOCK_SIZE
read1 := fun pid _ j => pid * X_row_stride + j.val
read2 := fun _ _ j => j.val
read3 := fun _ _ j => j.val
write1 := fun pid _ j => pid * Y_row_stride + j.val
write2 := fun pid _ _ => pid
write3 := fun pid _ _ => pid
mask := fun _ _ j => j.val < n_cols
writeMask2 := fun _ _ j => j.val = 0
writeMask3 := fun _ _ j => j.val = 0
Referenced model/specification definitions: Python source |
|---|
fast_rms_layernormrms_layernorm_forward_correctnessMathematical execution | Configured model / stage Separate ordinary and Gemma forward branches with Y/inverse-variance outputs. The Python backward kernel and wrapper composition are outside these headlines. | Statement and preconditionsspecification rms_layernorm_forward_correctness
(Y X W r : RegionName)
(Y_row_stride X_row_stride W_row_stride r_row_stride n_cols BLOCK_SIZE : Nat)
(eps : ℝ) (hYr : Y ≠ r) (hB : 0 < BLOCK_SIZE) :
rmsLayernormFwdIO Y X W r Y_row_stride X_row_stride W_row_stride
r_row_stride n_cols eps BLOCK_SIZE ⊨
fun _ _ xs ws =>
(fun i => rmsFwdYSpec n_cols BLOCK_SIZE eps xs ws i,
fun _ => rmsFwdInvVarSpec n_cols BLOCK_SIZE eps xs)
IO definition: rmsLayernormFwdIOdef rmsLayernormFwdIO (Y X W r : RegionName)
(Y_row_stride X_row_stride W_row_stride r_row_stride n_cols : Nat)
(eps : ℝ) (BLOCK_SIZE : Nat) : Masked2DKernelIO₂ₓ₂ where
kernel := rms_layernorm_forward Y X W r Y_row_stride X_row_stride
W_row_stride r_row_stride n_cols eps BLOCK_SIZE
in1 := X
in2 := W
out1 := Y
out2 := r
B := BLOCK_SIZE
read1 := fun row_idx _ j => row_idx * X_row_stride + j.val
read2 := fun _ _ j => j.val * W_row_stride
write1 := fun row_idx _ j => row_idx * Y_row_stride + j.val
write2 := fun row_idx _ _ => row_idx * r_row_stride
mask := fun _ _ j => j.val < n_cols
writeMask2 := fun _ _ j => j.val = 0
Referenced model/specification definitions: Python source |
|---|
fast_rms_layernormgemma_rms_layernorm_forward_correctnessMathematical execution | Configured model / stage Separate ordinary and Gemma forward branches with Y/inverse-variance outputs. The Python backward kernel and wrapper composition are outside these headlines. | Statement and preconditionsspecification gemma_rms_layernorm_forward_correctness
(Y X W r : RegionName)
(Y_row_stride X_row_stride r_row_stride n_cols BLOCK_SIZE : Nat)
(eps : ℝ) (hYr : Y ≠ r) (hB : 0 < BLOCK_SIZE) :
gemmaRmsLayernormFwdIO Y X W r Y_row_stride X_row_stride r_row_stride
n_cols eps BLOCK_SIZE ⊨
fun _ _ xs ws =>
(fun i => gemmaRmsFwdYSpec n_cols BLOCK_SIZE eps xs ws i,
fun _ => rmsFwdInvVarSpec n_cols BLOCK_SIZE eps xs)
IO definition: gemmaRmsLayernormFwdIOdef gemmaRmsLayernormFwdIO (Y X W r : RegionName)
(Y_row_stride X_row_stride r_row_stride n_cols : Nat)
(eps : ℝ) (BLOCK_SIZE : Nat) : Masked2DKernelIO₂ₓ₂ where
kernel := gemma_rms_layernorm_forward Y X W r Y_row_stride X_row_stride
r_row_stride n_cols eps BLOCK_SIZE
in1 := X
in2 := W
out1 := Y
out2 := r
B := BLOCK_SIZE
read1 := fun row_idx _ j => row_idx * X_row_stride + j.val
read2 := fun _ _ j => j.val
write1 := fun row_idx _ j => row_idx * Y_row_stride + j.val
write2 := fun row_idx _ _ => row_idx * r_row_stride
mask := fun _ _ j => j.val < n_cols
writeMask2 := fun _ _ j => j.val = 0
Referenced model/specification definitions: Python source |
|---|
fast_rope_embeddingrope_embedding_output_summary_generalMathematical execution | Configured model / stage Separate first/second-half rotary slices. Projection facts for complete/backward surfaces are not value proofs or a composed in-place two-half contract. | Statement and preconditionsspecification rope_embedding_output_summary_general
(Q K QGrad KGrad cos sin : RegionName)
(Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
BLOCK_SIZE : Nat)
(sQ sK : BlockState)
(hQFirstInj : Function.Injective
(fun i : Fin BLOCK_SIZE => qFirstOffset sQ Q_row_stride head_dim i))
(hQSecondInj : Function.Injective
(fun i : Fin BLOCK_SIZE => qSecondOffset sQ Q_row_stride head_dim i))
(hKFirstInj : Function.Injective
(fun i : Fin BLOCK_SIZE => qFirstOffset sK Q_row_stride head_dim i))
(hKSecondInj : Function.Injective
(fun i : Fin BLOCK_SIZE => qSecondOffset sK Q_row_stride head_dim i)) :
(∃ alg, (rope_embedding_surface Q cos sin Q_row_stride cos_row_stride
sin_row_stride seqlen head_dim n_heads BLOCK_SIZE
Bool.false).toAlgorithm? = Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := rope_embedding_q_first_half Q cos sin Q_row_stride
cos_row_stride sin_row_stride seqlen head_dim n_heads BLOCK_SIZE)
(initialState := sQ)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => active sQ head_dim n_heads BLOCK_SIZE i)
(fun i => (Q, qFirstOffset sQ Q_row_stride head_dim i)))
(expected := fun i =>
ropeFirstSpec sQ Q cos sin Q_row_stride cos_row_stride sin_row_stride
seqlen head_dim BLOCK_SIZE i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := rope_embedding_q_second_half Q cos sin Q_row_stride
cos_row_stride sin_row_stride seqlen head_dim n_heads BLOCK_SIZE)
(initialState := sQ)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => active sQ head_dim n_heads BLOCK_SIZE i)
(fun i => (Q, qSecondOffset sQ Q_row_stride head_dim i)))
(expected := fun i =>
ropeSecondSpec sQ Q cos sin Q_row_stride cos_row_stride sin_row_stride
seqlen head_dim BLOCK_SIZE i)) ∧
(∃ alg, (rope_embedding_surface K cos sin Q_row_stride cos_row_stride
sin_row_stride seqlen head_dim n_heads BLOCK_SIZE
Bool.false).toAlgorithm? = Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := rope_embedding_q_first_half K cos sin Q_row_stride
cos_row_stride sin_row_stride seqlen head_dim n_heads BLOCK_SIZE)
(initialState := sK)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => active sK head_dim n_heads BLOCK_SIZE i)
(fun i => (K, qFirstOffset sK Q_row_stride head_dim i)))
(expected := fun i =>
ropeFirstSpec sK K cos sin Q_row_stride cos_row_stride sin_row_stride
seqlen head_dim BLOCK_SIZE i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := rope_embedding_q_second_half K cos sin Q_row_stride
cos_row_stride sin_row_stride seqlen head_dim n_heads BLOCK_SIZE)
(initialState := sK)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => active sK head_dim n_heads BLOCK_SIZE i)
(fun i => (K, qSecondOffset sK Q_row_stride head_dim i)))
(expected := fun i =>
ropeSecondSpec sK K cos sin Q_row_stride cos_row_stride sin_row_stride
seqlen head_dim BLOCK_SIZE i)) ∧
(∃ alg, (rope_embedding_surface QGrad cos sin Q_row_stride cos_row_stride
sin_row_stride seqlen head_dim n_heads BLOCK_SIZE
Bool.true).toAlgorithm? = Except.ok alg) ∧
(∃ alg, (rope_embedding_surface KGrad cos sin Q_row_stride cos_row_stride
sin_row_stride seqlen head_dim n_heads BLOCK_SIZE
Bool.true).toAlgorithm? = Except.ok alg)
Referenced model/specification definitions: Python source |
|---|
fast_rope_embeddingfast_rope_embedding_io_correctnessMathematical execution | Configured model / stage Separate first/second-half rotary slices. Projection facts for complete/backward surfaces are not value proofs or a composed in-place two-half contract. | Statement and preconditionsspecification fast_rope_embedding_io_correctness (Q cos sin : RegionName)
(Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
BLOCK_SIZE : Nat) :
(ropeFirstIO Q cos sin Q_row_stride cos_row_stride sin_row_stride seqlen
head_dim n_heads BLOCK_SIZE
⊨ fun _p₀ _p₁ q1 q2 c1 s1 i => ropeFirstSpecOf q1 q2 c1 s1 i) ∧
(ropeSecondIO Q cos sin Q_row_stride cos_row_stride sin_row_stride seqlen
head_dim n_heads BLOCK_SIZE
⊨ fun _p₀ _p₁ q1 q2 c1 s1 i => ropeSecondSpecOf q1 q2 c1 s1 i)
IO definition: ropeFirstIOdef ropeFirstIO (Q cos sin : RegionName)
(Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
BLOCK_SIZE : Nat) : InPlaceMaskedTileKernelIO where
kernel := rope_embedding_q_first_half Q cos sin Q_row_stride cos_row_stride
sin_row_stride seqlen head_dim n_heads BLOCK_SIZE
main := Q
aux1 := cos
aux2 := sin
shape := [BLOCK_SIZE]
readMain1 := fun p₀ p₁ i => p₀ * Q_row_stride + p₁ * 4 * head_dim + i.1.val
readMain2 := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * 4 * head_dim + i.1.val + head_dim / 2
readAux1 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * cos_row_stride + i.1.val
readAux2 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * sin_row_stride + i.1.val
write := fun p₀ p₁ i => p₀ * Q_row_stride + p₁ * 4 * head_dim + i.1.val
mask := fun _p₀ _p₁ i => i.1.val < head_dim / 2
writeMask := fun _p₀ p₁ i => i.1.val < head_dim / 2 ∧ p₁ * 4 < n_heads
IO definition: ropeSecondIOdef ropeSecondIO (Q cos sin : RegionName)
(Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
BLOCK_SIZE : Nat) : InPlaceMaskedTileKernelIO where
kernel := rope_embedding_q_second_half Q cos sin Q_row_stride cos_row_stride
sin_row_stride seqlen head_dim n_heads BLOCK_SIZE
main := Q
aux1 := cos
aux2 := sin
shape := [BLOCK_SIZE]
readMain1 := fun p₀ p₁ i => p₀ * Q_row_stride + p₁ * 4 * head_dim + i.1.val
readMain2 := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * 4 * head_dim + i.1.val + head_dim / 2
readAux1 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * cos_row_stride + i.1.val
readAux2 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * sin_row_stride + i.1.val
write := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * 4 * head_dim + i.1.val + head_dim / 2
mask := fun _p₀ _p₁ i => i.1.val < head_dim / 2
writeMask := fun _p₀ p₁ i => i.1.val < head_dim / 2 ∧ p₁ * 4 < n_heads
Referenced model/specification definitions: Python source |
|---|
fifth_order_sph_harmonicsfifth_order_fwd_correctnessMathematical execution | Configured model / stage Forward real-valued harmonic outputs under complete XYZ-coordinate and output-bound premises. The backward kernel and hardware transcendental/numeric behavior are external. | Statement and preconditionsspecification fifth_order_fwd_correctness
(coord_ptr output_ptr : RegionName)
(block_size coord_numel output_numel col_offset output_stride : Nat)
(hStride : 10 < output_stride)
(hCover : ∀ (pid₀ : Nat) (j : Fin block_size),
j.val * output_stride + block_size * output_stride * pid₀ + col_offset
< output_numel →
j.val * 3 + block_size * 3 * pid₀ + 2 < coord_numel) :
fifthOrderFwdIO coord_ptr output_ptr block_size coord_numel output_numel
col_offset output_stride ⊨
fun _ _ xs o j =>
sphY o (xs ⟨0, by show 0 < 3; decide⟩ j) (xs ⟨1, by show 1 < 3; decide⟩ j)
(xs ⟨2, by show 2 < 3; decide⟩ j)
IO definition: fifthOrderFwdIOdef fifthOrderFwdIO (coord_ptr output_ptr : RegionName)
(block_size coord_numel output_numel col_offset output_stride : Nat) :
GroupedMasked2DKernelIO where
kernel := fifth_order_fwd_surface coord_ptr output_ptr block_size coord_numel
output_numel col_offset output_stride
projection := by
obtain ⟨alg, h⟩ := fifth_order_fwd_surface_toAlgorithm_supported coord_ptr
output_ptr block_size coord_numel output_numel col_offset output_stride
simp only [ComputeKernel.toAlgKernel, h]
nIn := 3
nOut := 11
bufs := [coord_ptr, output_ptr]
inp := fun _ => coord_ptr
out := fun _ => output_ptr
B := block_size
read := fun i pid₀ _ j => sphInWin block_size pid₀ i j
readMask := fun i pid₀ _ j => sphInWin block_size pid₀ i j < coord_numel
write := fun o pid₀ _ j => sphOutWin block_size col_offset output_stride pid₀ o j
writeMask := fun o pid₀ _ j =>
sphOutWin block_size col_offset output_stride pid₀ o j < output_numel
Referenced model/specification definitions: Python source |
|---|
flash_attnflash_attn_genuine_output_compute_correct_generalMathematical execution | Configured model / stage Forward causal/non-causal value facts at the displayed contiguous-layout specialization. The IO headline fixes non-causal mode; O and L are separate outputs and identity-fp16 premises must be read where present. | Statement and preconditionsspecification flash_attn_genuine_output_compute_correct_general
(Q K V L O : RegionName) (s : BlockState) (IS_CAUSAL : Bool)
(sm_scale : ℝ) (stride_q_head SEQLEN BLOCK_M DIM BLOCK_N : Nat)
(sqbs skbs svbs sobs sosl sod BS HEAD : Nat)
(hDIM : 0 < DIM) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M) (hBMlen : 1 < [BLOCK_M].length.succ)
(hdvd : BLOCK_N ∣ SEQLEN) (hSEQ : 0 < SEQLEN)
(hHi : flashHiG s IS_CAUSAL SEQLEN BLOCK_M = SEQLEN)
(hOL : O ≠ L) (hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := flash_attn_fwd_kernel_surface Q K V L O sm_scale
sqbs stride_q_head DIM 1 skbs stride_q_head DIM 1 svbs stride_q_head DIM 1
sobs stride_q_head DIM 1 BS HEAD SEQLEN BLOCK_M DIM BLOCK_N IS_CAUSAL)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, DIM] =>
some (O, outOffset s stride_q_head DIM 1 BLOCK_M idx))
(expected := fun idx : TileIndex [BLOCK_M, DIM] =>
MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (if IS_CAUSAL then
flashAttnOValueSpecCausal s Q K V sm_scale stride_q_head DIM SEQLEN BLOCK_M idx
else
flashAttnOValueSpec s Q K V sm_scale stride_q_head DIM SEQLEN BLOCK_M idx))))
Referenced model/specification definitions: Python source |
|---|
flash_attnflash_attn_genuine_l_compute_correct_generalMathematical execution | Configured model / stage Forward causal/non-causal value facts at the displayed contiguous-layout specialization. The IO headline fixes non-causal mode; O and L are separate outputs and identity-fp16 premises must be read where present. | Statement and preconditionsspecification flash_attn_genuine_l_compute_correct_general
(Q K V L O : RegionName) (s : BlockState) (IS_CAUSAL : Bool)
(sm_scale : ℝ) (stride_q_head SEQLEN BLOCK_M DIM BLOCK_N : Nat)
(sqbs skbs svbs sobs sosl sod BS HEAD : Nat)
(hDIM : 0 < DIM) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M) (hBMlen : 1 < [BLOCK_M].length.succ)
(hdvd : BLOCK_N ∣ SEQLEN) (hSEQ : 0 < SEQLEN)
(hHi : flashHiG s IS_CAUSAL SEQLEN BLOCK_M = SEQLEN)
(hOL : O ≠ L) (hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := flash_attn_fwd_kernel_surface Q K V L O sm_scale
sqbs stride_q_head DIM 1 skbs stride_q_head DIM 1 svbs stride_q_head DIM 1
sobs stride_q_head DIM 1 BS HEAD SEQLEN BLOCK_M DIM BLOCK_N IS_CAUSAL)
(initialState := s)
(write := fun i : Fin BLOCK_M => some (L, lOffset s SEQLEN BLOCK_M i))
(expected := fun i : Fin BLOCK_M =>
Real.log
(((flashKeysUpto (qTile s Q stride_q_head DIM BLOCK_M) (kTile s K stride_q_head DIM SEQLEN)
(vTile s V stride_q_head DIM SEQLEN) (sm_scale * log2e) IS_CAUSAL (s.pids 0 * BLOCK_M) SEQLEN i
⟨0, hDIM⟩).map (fun p => pow2 p.1)).sum) / Real.log 2)
Referenced model/specification definitions: Python source |
|---|
flash_attnflash_attn_python_case1_genuine_compute_correct_generalMathematical execution | Configured model / stage Forward causal/non-causal value facts at the displayed contiguous-layout specialization. The IO headline fixes non-causal mode; O and L are separate outputs and identity-fp16 premises must be read where present. | Statement and preconditionsspecification flash_attn_python_case1_genuine_compute_correct_general
(Q K V L O : RegionName) (s : BlockState)
(sm_scale : ℝ) (stride_q_head SEQLEN BLOCK_M DIM BLOCK_N : Nat)
(sqbs skbs svbs sobs sosl sod BS HEAD : Nat)
(hDIM : 0 < DIM) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M) (hBMlen : 1 < [BLOCK_M].length.succ)
(hdvd : BLOCK_N ∣ SEQLEN) (hSEQ : 0 < SEQLEN)
(hAlign : (s.pids 0 + 1) * BLOCK_M = SEQLEN)
(hOL : O ≠ L) (hundef : ∀ rg o, s.undef rg o = 0) :
(ComputeCorrect.Realizes_without_Rounding
(kernel := flash_attn_fwd_kernel_surface Q K V L O sm_scale
sqbs stride_q_head DIM 1 skbs stride_q_head DIM 1 svbs stride_q_head DIM 1
sobs stride_q_head DIM 1 BS HEAD SEQLEN BLOCK_M DIM BLOCK_N Bool.true)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, DIM] => some (O, outOffset s stride_q_head DIM 1 BLOCK_M idx))
(expected := fun idx : TileIndex [BLOCK_M, DIM] =>
MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (flashAttnOValueSpecCausal s Q K V sm_scale stride_q_head DIM SEQLEN BLOCK_M idx))))) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := flash_attn_fwd_kernel_surface Q K V L O sm_scale
sqbs stride_q_head DIM 1 skbs stride_q_head DIM 1 svbs stride_q_head DIM 1
sobs stride_q_head DIM 1 BS HEAD SEQLEN BLOCK_M DIM BLOCK_N Bool.true)
(initialState := s)
(write := fun i : Fin BLOCK_M => some (L, lOffset s SEQLEN BLOCK_M i))
(expected := fun i : Fin BLOCK_M =>
Real.log
(((flashKeysUpto (qTile s Q stride_q_head DIM BLOCK_M) (kTile s K stride_q_head DIM SEQLEN)
(vTile s V stride_q_head DIM SEQLEN) (sm_scale * log2e) Bool.true (s.pids 0 * BLOCK_M) SEQLEN i
⟨0, hDIM⟩).map (fun p => pow2 p.1)).sum) / Real.log 2))
Referenced model/specification definitions: Python source |
|---|
flash_attnflash_attn_python_case2_genuine_compute_correct_generalMathematical execution | Configured model / stage Forward causal/non-causal value facts at the displayed contiguous-layout specialization. The IO headline fixes non-causal mode; O and L are separate outputs and identity-fp16 premises must be read where present. | Statement and preconditionsspecification flash_attn_python_case2_genuine_compute_correct_general
(Q K V L O : RegionName) (s : BlockState)
(sm_scale : ℝ) (stride_q_head SEQLEN BLOCK_M DIM BLOCK_N : Nat)
(sqbs skbs svbs sobs sosl sod BS HEAD : Nat)
(hDIM : 0 < DIM) (hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M) (hBMlen : 1 < [BLOCK_M].length.succ)
(hdvd : BLOCK_N ∣ SEQLEN) (hSEQ : 0 < SEQLEN)
(hOL : O ≠ L) (hundef : ∀ rg o, s.undef rg o = 0) :
(ComputeCorrect.Realizes_without_Rounding
(kernel := flash_attn_fwd_kernel_surface Q K V L O sm_scale
sqbs stride_q_head DIM 1 skbs stride_q_head DIM 1 svbs stride_q_head DIM 1
sobs stride_q_head DIM 1 BS HEAD SEQLEN BLOCK_M DIM BLOCK_N Bool.false)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, DIM] => some (O, outOffset s stride_q_head DIM 1 BLOCK_M idx))
(expected := fun idx : TileIndex [BLOCK_M, DIM] =>
MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (flashAttnOValueSpec s Q K V sm_scale stride_q_head DIM SEQLEN BLOCK_M idx))))) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := flash_attn_fwd_kernel_surface Q K V L O sm_scale
sqbs stride_q_head DIM 1 skbs stride_q_head DIM 1 svbs stride_q_head DIM 1
sobs stride_q_head DIM 1 BS HEAD SEQLEN BLOCK_M DIM BLOCK_N Bool.false)
(initialState := s)
(write := fun i : Fin BLOCK_M => some (L, lOffset s SEQLEN BLOCK_M i))
(expected := fun i : Fin BLOCK_M =>
Real.log
(((flashKeysUpto (qTile s Q stride_q_head DIM BLOCK_M) (kTile s K stride_q_head DIM SEQLEN)
(vTile s V stride_q_head DIM SEQLEN) (sm_scale * log2e) Bool.false (s.pids 0 * BLOCK_M) SEQLEN i
⟨0, hDIM⟩).map (fun p => pow2 p.1)).sum) / Real.log 2))
Referenced model/specification definitions: Python source |
|---|
flash_attnflash_attn_io_correctnessAbstract cast/store rounding | Configured model / stage Forward causal/non-causal value facts at the displayed contiguous-layout specialization. The IO headline fixes non-causal mode; O and L are separate outputs and identity-fp16 premises must be read where present. | Statement and preconditionsspecification flash_attn_io_correctness (R : RoundingModel) (hfp16 : R.round .fp16 = id)
(Q K V L O : RegionName) (sm_scale : ℝ)
(sqbs skbs svbs sobs BS HEAD SEQLEN BLOCK_M DIM BLOCK_N stride_q_head : Nat)
(hDIM : 0 < DIM) (hBN : 0 < BLOCK_N) (hSEQ : 0 < SEQLEN) (hdvd : BLOCK_N ∣ SEQLEN)
(hOL : O ≠ L) :
flashAttnIO Q K V L O sm_scale sqbs skbs svbs sobs BS HEAD SEQLEN BLOCK_M DIM BLOCK_N
stride_q_head ⊨[R]
fun _ _ _ xs ys zs =>
(fun j => flashIOOutSpec BLOCK_M DIM BLOCK_N SEQLEN (SEQLEN / BLOCK_N)
(Nat.div_pos (Nat.le_of_dvd hSEQ hdvd) hBN) (Nat.div_mul_cancel hdvd) hBN
sm_scale xs ys zs j,
fun i => flashIOLSpec BLOCK_M DIM BLOCK_N SEQLEN (SEQLEN / BLOCK_N) hDIM
(Nat.div_pos (Nat.le_of_dvd hSEQ hdvd) hBN) (Nat.div_mul_cancel hdvd) hBN
sm_scale xs ys zs i)
IO definition: flashAttnIOdef flashAttnIO (Q K V L O : RegionName) (sm_scale : ℝ)
(sqbs skbs svbs sobs BS HEAD SEQLEN BLOCK_M DIM BLOCK_N stride_q_head : Nat) :
StreamMasked3DKernelIO₃ₓ₂ where
kernel := flash_attn_fwd_kernel_surface Q K V L O sm_scale
sqbs stride_q_head DIM 1 skbs stride_q_head DIM 1 svbs stride_q_head DIM 1
sobs stride_q_head DIM 1 BS HEAD SEQLEN BLOCK_M DIM BLOCK_N Bool.false
inp1 := Q
inp2 := K
inp3 := V
out1 := O
out2 := L
T := SEQLEN / BLOCK_N
B1 := BLOCK_M * DIM
B2 := DIM * BLOCK_N
B3 := BLOCK_N * DIM
C1 := BLOCK_M * DIM
C2 := BLOCK_M
out1DType := .fp16
out2DType := .real
read1 := fun p₀ p₁ _ _ j =>
p₁ * stride_q_head + (p₀ * BLOCK_M + j.val / DIM) * DIM + (j.val % DIM) * 1
read2 := fun _ p₁ _ t j =>
p₁ * stride_q_head + (j.val / BLOCK_N) * 1 + (t.val * BLOCK_N + j.val % BLOCK_N) * DIM
read3 := fun _ p₁ _ t j =>
p₁ * stride_q_head + (t.val * BLOCK_N + j.val / DIM) * DIM + (j.val % DIM) * 1
write1 := fun p₀ p₁ _ j =>
p₁ * stride_q_head + (p₀ * BLOCK_M + j.val / DIM) * DIM + (j.val % DIM) * 1
write2 := fun p₀ p₁ _ i => p₁ * SEQLEN + (p₀ * BLOCK_M + i.val)
mask1 := fun _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ => True
mask3 := fun _ _ _ _ _ => True
writeMask1 := fun _ _ _ _ => True
writeMask2 := fun _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
flash_decode2_llamaflash_decode2_llama_normalization_output_summary_generalMathematical execution | Configured model / stage Stage-2 normalization and running-max facts include both producer and step kernels; each conjunct names its own target. Inputs from stage 1 remain supplied. | Statement and preconditionsspecification flash_decode2_llama_normalization_output_summary_general
(B_Seqlen : Region .nat) (Mid_O Mid_O_LogExpSum Acc SumExp O : RegionName)
(stride_mid_ob stride_mid_oh stride_mid_os stride_mid_od
stride_mid_o_eb stride_mid_o_eh stride_mid_o_es stride_obs stride_oh stride_od
BLOCK_SEQ BLOCK_DMODEL
stride_acc_b stride_acc_h stride_acc_d stride_sum_b stride_sum_h : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_DMODEL => outOffset s stride_obs stride_oh stride_od i)) :
(∃ alg, (flash_decode2_llama_surface B_Seqlen Mid_O Mid_O_LogExpSum O
stride_mid_ob stride_mid_oh stride_mid_os stride_mid_od stride_mid_o_eb
stride_mid_o_eh stride_mid_o_es stride_obs stride_oh stride_od BLOCK_SEQ
BLOCK_DMODEL).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := flash_decode2_llama_normalization_store_kernel Acc SumExp O
stride_acc_b stride_acc_h stride_acc_d stride_sum_b stride_sum_h
stride_obs stride_oh stride_od BLOCK_DMODEL)
(initialState := s)
(write := fun i : Fin BLOCK_DMODEL =>
some (O, outOffset s stride_obs stride_oh stride_od i))
(expected := fun i : Fin BLOCK_DMODEL =>
normalizedStoreValue s Acc SumExp stride_acc_b stride_acc_h
stride_acc_d stride_sum_b stride_sum_h i)
Referenced model/specification definitions: Python source |
|---|
flash_decode2_llamaflash_decode2_llama_running_max_output_summary_generalMathematical execution | Configured model / stage Stage-2 normalization and running-max facts include both producer and step kernels; each conjunct names its own target. Inputs from stage 1 remain supplied. | Statement and preconditionsspecification flash_decode2_llama_running_max_output_summary_general
(B_Seqlen : Region .nat)
(Mid_O Mid_O_LogExpSum MaxLogic NewMaxLogic O : RegionName)
(stride_mid_ob stride_mid_oh stride_mid_os stride_mid_od
stride_mid_o_eb' stride_mid_o_eh' stride_mid_o_es stride_obs stride_oh stride_od
BLOCK_SEQ BLOCK_DMODEL
block_seq_n stride_mid_o_eb stride_mid_o_eh stride_logic_b stride_logic_h : Nat)
(s : BlockState) :
(∃ alg, (flash_decode2_llama_surface B_Seqlen Mid_O Mid_O_LogExpSum O
stride_mid_ob stride_mid_oh stride_mid_os stride_mid_od stride_mid_o_eb'
stride_mid_o_eh' stride_mid_o_es stride_obs stride_oh stride_od BLOCK_SEQ
BLOCK_DMODEL).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := flash_decode2_llama_running_max_step_kernel Mid_O_LogExpSum
MaxLogic NewMaxLogic block_seq_n stride_mid_o_eb stride_mid_o_eh
stride_logic_b stride_logic_h)
(initialState := s)
(write := ComputeCorrect.WriteMap.scalar NewMaxLogic
(sumExpOffset s stride_logic_b stride_logic_h))
(expected := fun _ : PUnit =>
runningMaxStepValue s Mid_O_LogExpSum MaxLogic block_seq_n
stride_mid_o_eb stride_mid_o_eh stride_logic_b stride_logic_h)
Referenced model/specification definitions: Python source |
|---|
flash_decode2_llamaflash_decode2_llama_final_store_io_correctnessMathematical execution | Precomputed-input slice Copies supplied Final values into O; neither the running maximum nor normalization is computed by this IO kernel. | Statement and preconditionsspecification flash_decode2_llama_final_store_io_correctness (Final O : RegionName)
(stride_final_b stride_final_h stride_final_d stride_obs stride_oh
stride_od BLOCK_DMODEL : Nat)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun i : Fin BLOCK_DMODEL =>
p₀ * stride_obs + p₁ * stride_oh + i.val * stride_od)) :
finalStoreIO Final O stride_final_b stride_final_h stride_final_d stride_obs
stride_oh stride_od BLOCK_DMODEL
⊨ fun _p₀ _p₁ xs i => xs i
IO definition: finalStoreIOdef finalStoreIO (Final O : RegionName)
(stride_final_b stride_final_h stride_final_d stride_obs stride_oh
stride_od BLOCK_DMODEL : Nat) : Masked3DTileKernelIO₁ where
kernel := flash_decode2_llama_final_store_slice Final O stride_final_b
stride_final_h stride_final_d stride_obs stride_oh stride_od BLOCK_DMODEL
inp := Final
out := O
shape := [BLOCK_DMODEL]
read := fun p₀ p₁ _p₂ i =>
p₀ * stride_final_b + p₁ * stride_final_h + i.1.val * stride_final_d
write := fun p₀ p₁ _p₂ i =>
p₀ * stride_obs + p₁ * stride_oh + i.1.val * stride_od
mask := fun _p₀ _p₁ _p₂ _ => True
Referenced model/specification definitions: Python source |
|---|
flash_decode2_llamaflash_decode2_llama_final_store_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding version of the supplied-Final output-store slice. | Statement and preconditionsspecification flash_decode2_llama_final_store_io_correctnessR (R : RoundingModel)
(Final O : RegionName)
(stride_final_b stride_final_h stride_final_d stride_obs stride_oh
stride_od BLOCK_DMODEL : Nat)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun i : Fin BLOCK_DMODEL =>
p₀ * stride_obs + p₁ * stride_oh + i.val * stride_od)) :
finalStoreIO Final O stride_final_b stride_final_h stride_final_d stride_obs
stride_oh stride_od BLOCK_DMODEL
⊨[R, FloatDType.real] fun _p₀ _p₁ xs i => xs i
IO definition: finalStoreIOdef finalStoreIO (Final O : RegionName)
(stride_final_b stride_final_h stride_final_d stride_obs stride_oh
stride_od BLOCK_DMODEL : Nat) : Masked3DTileKernelIO₁ where
kernel := flash_decode2_llama_final_store_slice Final O stride_final_b
stride_final_h stride_final_d stride_obs stride_oh stride_od BLOCK_DMODEL
inp := Final
out := O
shape := [BLOCK_DMODEL]
read := fun p₀ p₁ _p₂ i =>
p₀ * stride_final_b + p₁ * stride_final_h + i.1.val * stride_final_d
write := fun p₀ p₁ _p₂ i =>
p₀ * stride_obs + p₁ * stride_oh + i.1.val * stride_od
mask := fun _p₀ _p₁ _p₂ _ => True
Referenced model/specification definitions: Python source |
|---|
flash_decode2_phiflash_decode2_phi_masked_accumulator_output_summary_generalMathematical execution | Configured model / stage Stage-2 accumulator, running-max, and normalization facts, with masked dimensions and step premises. Stage-1 inputs are supplied and the summary conjuncts have different kernels. | Statement and preconditionsspecification flash_decode2_phi_masked_accumulator_output_summary_general
(B_Seqlen : Region .nat)
(Mid_O Mid_O_LogExpSum AccIn SumExpIn MaxLogic NewMaxLogic AccOut SumExpOut
Out : RegionName)
(stride_mid_ob stride_mid_oh stride_mid_os stride_mid_od
stride_mid_o_eb stride_mid_o_eh stride_mid_o_es stride_obs stride_oh stride_od
head_dim BLOCK_SEQ BLOCK_DMODEL
block_seq_n stride_acc_b stride_acc_h stride_acc_d
stride_sum_b stride_sum_h stride_logic_b stride_logic_h : Nat)
(s : BlockState)
(hAccOutInj : Function.Injective
(fun i : Fin BLOCK_DMODEL =>
accOffset s stride_acc_b stride_acc_h stride_acc_d i)) :
(∃ alg, (flash_decode2_phi_surface B_Seqlen Mid_O Mid_O_LogExpSum Out
stride_mid_ob stride_mid_oh stride_mid_os stride_mid_od stride_mid_o_eb
stride_mid_o_eh stride_mid_o_es stride_obs stride_oh stride_od head_dim
BLOCK_SEQ BLOCK_DMODEL).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := flash_decode2_phi_accumulator_step_kernel Mid_O Mid_O_LogExpSum
AccIn MaxLogic NewMaxLogic AccOut block_seq_n head_dim stride_mid_ob
stride_mid_oh stride_mid_os stride_mid_od stride_mid_o_eb stride_mid_o_eh
stride_acc_b stride_acc_h stride_acc_d stride_logic_b stride_logic_h
BLOCK_DMODEL)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_DMODEL => active s head_dim i)
(fun i => (AccOut, accOffset s stride_acc_b stride_acc_h stride_acc_d i)))
(expected := fun i : Fin BLOCK_DMODEL =>
accumulatorStepValue s Mid_O Mid_O_LogExpSum AccIn MaxLogic NewMaxLogic
block_seq_n head_dim stride_mid_ob stride_mid_oh stride_mid_os
stride_mid_od stride_mid_o_eb stride_mid_o_eh stride_acc_b stride_acc_h
stride_acc_d stride_logic_b stride_logic_h i) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := flash_decode2_phi_sum_exp_step_kernel Mid_O_LogExpSum SumExpIn
MaxLogic NewMaxLogic SumExpOut block_seq_n stride_mid_o_eb stride_mid_o_eh
stride_sum_b stride_sum_h stride_logic_b stride_logic_h)
(initialState := s)
(write := ComputeCorrect.WriteMap.scalar SumExpOut
(sumExpOffset s stride_sum_b stride_sum_h))
(expected := fun _ : PUnit =>
sumExpStepValue s Mid_O_LogExpSum SumExpIn MaxLogic NewMaxLogic
block_seq_n stride_mid_o_eb stride_mid_o_eh stride_sum_b stride_sum_h
stride_logic_b stride_logic_h)
Referenced model/specification definitions: Python source |
|---|
flash_decode2_phiflash_decode2_phi_running_max_output_summary_generalMathematical execution | Configured model / stage Stage-2 accumulator, running-max, and normalization facts, with masked dimensions and step premises. Stage-1 inputs are supplied and the summary conjuncts have different kernels. | Statement and preconditionsspecification flash_decode2_phi_running_max_output_summary_general
(B_Seqlen : Region .nat) (Mid_O Mid_O_LogExpSum Out : RegionName)
(stride_mid_ob stride_mid_oh stride_mid_os stride_mid_od
stride_mid_o_eb stride_mid_o_eh stride_mid_o_es stride_obs stride_oh stride_od
head_dim BLOCK_SEQ BLOCK_DMODEL : Nat)
(s : BlockState) (block_n_size : Nat) :
(∃ alg, (flash_decode2_phi_surface B_Seqlen Mid_O Mid_O_LogExpSum Out
stride_mid_ob stride_mid_oh stride_mid_os stride_mid_od stride_mid_o_eb
stride_mid_o_eh stride_mid_o_es stride_obs stride_oh stride_od head_dim
BLOCK_SEQ BLOCK_DMODEL).toAlgorithm? = Except.ok alg) ∧
runningMaxAfter s Mid_O_LogExpSum stride_mid_o_eb stride_mid_o_eh 0 = ⊥ ∧
(∀ k, k < block_n_size →
runningMaxAfter s Mid_O_LogExpSum stride_mid_o_eb stride_mid_o_eh (k + 1) =
runningMaxJoin
(some (s.readMem Mid_O_LogExpSum
(runningMaxLogicOffset s stride_mid_o_eb stride_mid_o_eh k)) :
WithBot ℝ)
(runningMaxAfter s Mid_O_LogExpSum stride_mid_o_eb stride_mid_o_eh k))
Referenced model/specification definitions: Python source |
|---|
flash_decode2_phiflash_decode2_phi_normalization_output_summary_generalMathematical execution | Configured model / stage Stage-2 accumulator, running-max, and normalization facts, with masked dimensions and step premises. Stage-1 inputs are supplied and the summary conjuncts have different kernels. | Statement and preconditionsspecification flash_decode2_phi_normalization_output_summary_general
(B_Seqlen : Region .nat) (Mid_O Mid_O_LogExpSum Acc SumExp Out : RegionName)
(stride_mid_ob stride_mid_oh stride_mid_os stride_mid_od
stride_mid_o_eb stride_mid_o_eh stride_mid_o_es stride_obs stride_oh stride_od
head_dim BLOCK_SEQ BLOCK_DMODEL
stride_acc_b stride_acc_h stride_acc_d stride_sum_b stride_sum_h : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_DMODEL => outOffset s stride_obs stride_oh stride_od i)) :
(∃ alg, (flash_decode2_phi_surface B_Seqlen Mid_O Mid_O_LogExpSum Out
stride_mid_ob stride_mid_oh stride_mid_os stride_mid_od stride_mid_o_eb
stride_mid_o_eh stride_mid_o_es stride_obs stride_oh stride_od head_dim
BLOCK_SEQ BLOCK_DMODEL).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := flash_decode2_phi_normalization_store_kernel Acc SumExp Out
head_dim stride_acc_b stride_acc_h stride_acc_d stride_sum_b stride_sum_h
stride_obs stride_oh stride_od BLOCK_DMODEL)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_DMODEL => active s head_dim i)
(fun i => (Out, outOffset s stride_obs stride_oh stride_od i)))
(expected := fun i : Fin BLOCK_DMODEL =>
normalizedStoreValue s Acc SumExp head_dim stride_acc_b stride_acc_h
stride_acc_d stride_sum_b stride_sum_h i)
Referenced model/specification definitions: Python source |
|---|
flash_decode2_phiflash_decode2_phi_final_store_io_correctnessMathematical execution | Precomputed-input slice Copies supplied Final values to active Out lanes; normalization and accumulator production are outside the IO kernel. | Statement and preconditionsspecification flash_decode2_phi_final_store_io_correctness (Final Out : RegionName)
(head_dim stride_final_b stride_final_h stride_final_d stride_obs stride_oh
stride_od BLOCK_DMODEL : Nat)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun i : Fin BLOCK_DMODEL =>
p₀ * stride_obs + p₁ * stride_oh + i.val * stride_od)) :
phiFinalStoreIO Final Out head_dim stride_final_b stride_final_h stride_final_d
stride_obs stride_oh stride_od BLOCK_DMODEL
⊨ fun _p₀ _p₁ xs i => xs i
IO definition: phiFinalStoreIOdef phiFinalStoreIO (Final Out : RegionName)
(head_dim stride_final_b stride_final_h stride_final_d stride_obs stride_oh
stride_od BLOCK_DMODEL : Nat) : Masked3DTileKernelIO₁ where
kernel := flash_decode2_phi_final_store_slice Final Out head_dim stride_final_b
stride_final_h stride_final_d stride_obs stride_oh stride_od BLOCK_DMODEL
inp := Final
out := Out
shape := [BLOCK_DMODEL]
read := fun p₀ p₁ _p₂ i =>
p₀ * stride_final_b + p₁ * stride_final_h + i.1.val * stride_final_d
write := fun p₀ p₁ _p₂ i =>
p₀ * stride_obs + p₁ * stride_oh + i.1.val * stride_od
mask := fun _p₀ _p₁ _p₂ i => i.1.val < head_dim
Referenced model/specification definitions: Python source |
|---|
flash_decode2_phiflash_decode2_phi_final_store_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding version of the supplied-Final masked output-store slice. | Statement and preconditionsspecification flash_decode2_phi_final_store_io_correctnessR (R : RoundingModel)
(Final Out : RegionName)
(head_dim stride_final_b stride_final_h stride_final_d stride_obs stride_oh
stride_od BLOCK_DMODEL : Nat)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun i : Fin BLOCK_DMODEL =>
p₀ * stride_obs + p₁ * stride_oh + i.val * stride_od)) :
phiFinalStoreIO Final Out head_dim stride_final_b stride_final_h stride_final_d
stride_obs stride_oh stride_od BLOCK_DMODEL
⊨[R, FloatDType.real] fun _p₀ _p₁ xs i => xs i
IO definition: phiFinalStoreIOdef phiFinalStoreIO (Final Out : RegionName)
(head_dim stride_final_b stride_final_h stride_final_d stride_obs stride_oh
stride_od BLOCK_DMODEL : Nat) : Masked3DTileKernelIO₁ where
kernel := flash_decode2_phi_final_store_slice Final Out head_dim stride_final_b
stride_final_h stride_final_d stride_obs stride_oh stride_od BLOCK_DMODEL
inp := Final
out := Out
shape := [BLOCK_DMODEL]
read := fun p₀ p₁ _p₂ i =>
p₀ * stride_final_b + p₁ * stride_final_h + i.1.val * stride_final_d
write := fun p₀ p₁ _p₂ i =>
p₀ * stride_obs + p₁ * stride_oh + i.1.val * stride_od
mask := fun _p₀ _p₁ _p₂ i => i.1.val < head_dim
Referenced model/specification definitions: Python source |
|---|
fused_activationfused_add_mul_activation_kernel_correctnessMathematical execution | Configured model / stage Selected per-program fused add/multiply and sigmoid-or-ReLU branch with the stated weight indexing and mask. Numeric operations have mathematical semantics. | Statement and preconditionsspecification fused_add_mul_activation_kernel_correctness
(x_ptr bias_ptr in_ptr : RegionName)
(num_weights xnumel BLOCK_SIZE : Nat)
(multiplier : ℝ) (ACTIVATION_SIGMOID : Bool) :
fusedActivationIO x_ptr bias_ptr in_ptr num_weights xnumel BLOCK_SIZE
multiplier ACTIVATION_SIGMOID
⊨ fun _ _ xs _ j =>
fusedActivationSpec ACTIVATION_SIGMOID
(xs (⟨0, by decide⟩ : Fin 3) j)
(xs (⟨1, by decide⟩ : Fin 3) j)
(xs (⟨2, by decide⟩ : Fin 3) j) multiplier
IO definition: fusedActivationIOdef fusedActivationIO
(x_ptr bias_ptr in_ptr : RegionName)
(num_weights xnumel BLOCK_SIZE : Nat)
(multiplier : ℝ) (ACTIVATION_SIGMOID : Bool) : GroupedMasked2DKernelIO where
kernel := fused_add_mul_activation_kernel x_ptr bias_ptr in_ptr num_weights
xnumel BLOCK_SIZE multiplier ACTIVATION_SIGMOID
nIn := 3
nOut := 1
bufs := [x_ptr, bias_ptr, in_ptr]
inp := fun i => match i with
| ⟨0, _⟩ => x_ptr
| ⟨1, _⟩ => bias_ptr
| ⟨_ + 2, _⟩ => in_ptr
out := fun _ => x_ptr
B := BLOCK_SIZE
read := fun i pid₀ _ j => match i with
| ⟨0, _⟩ => pid₀ * BLOCK_SIZE + j.val
| ⟨1, _⟩ => (pid₀ * BLOCK_SIZE + j.val) % num_weights
| ⟨_ + 2, _⟩ => pid₀ * BLOCK_SIZE + j.val
readMask := fun _ pid₀ _ j => pid₀ * BLOCK_SIZE + j.val < xnumel
write := fun _ pid₀ _ j => pid₀ * BLOCK_SIZE + j.val
writeMask := fun _ pid₀ _ j => pid₀ * BLOCK_SIZE + j.val < xnumel
Referenced model/specification definitions: Python source |
|---|
fused_layernorm_tritonfused_layernorm_triton_output_summary_generalMathematical execution | Configured model / stage Bundle of reduction, normalization, and linked-surface facts; the hypotheses connecting intermediate means/variances must hold. | Statement and preconditionsspecification fused_layernorm_triton_output_summary_general
(in_out_ptr0 in_ptr0 in_ptr1 in_ptr2 out_ptr0 out_ptr1 : RegionName)
(xnumel rnumel XBLOCK RBLOCK : Nat) (hLe : rnumel ≤ RBLOCK)
(s : BlockState)
(hMeanRstd : out_ptr0 ≠ in_out_ptr0) :
-- (1) the full faithful surface lowers to the algorithm layer
(∃ alg,
(fused_layernorm_triton_surface in_out_ptr0 in_ptr0 in_ptr1 in_ptr2
out_ptr0 out_ptr1 xnumel rnumel XBLOCK RBLOCK).toAlgorithm?
= Except.ok alg) ∧
-- (2) reduction phase: genuine mean and rstd, end-to-end from `in_ptr0`
((ComputeCorrect.Realizes_without_Rounding
(kernel := fused_layernorm_triton_reduce_slice in_out_ptr0 in_ptr0
out_ptr0 rnumel RBLOCK)
(initialState := s)
(write := fun _ : PUnit => some (out_ptr0, s.pids 0))
(expected := fun _ => rowMeanSpec s in_ptr0 rnumel)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_layernorm_triton_reduce_slice in_out_ptr0 in_ptr0
out_ptr0 rnumel RBLOCK)
(initialState := s)
(write := fun _ : PUnit => some (in_out_ptr0, s.pids 0))
(expected := fun _ => rowRstdSpec s in_ptr0 rnumel))) ∧
-- (3) normalize face: genuine `Y` chunk for every mean/rstd-genuine state
(∀ s' : BlockState,
s'.readMem out_ptr0 (s'.pids 0) = rowMeanSpec s' in_ptr0 rnumel →
s'.readMem in_out_ptr0 (s'.pids 0) = rowRstdSpec s' in_ptr0 rnumel →
ComputeCorrect.Realizes_without_Rounding
(kernel := fused_layernorm_triton_normalize_slice in_out_ptr0 in_ptr0
in_ptr1 in_ptr2 out_ptr0 out_ptr1 rnumel RBLOCK)
(initialState := s')
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin RBLOCK => activeLane s' rnumel RBLOCK i)
(fun i => (out_ptr1, yOffset s' rnumel RBLOCK i)))
(expected := fun i : Fin RBLOCK =>
rowYSpec s' in_ptr0 in_ptr1 in_ptr2 rnumel (rIndex s' RBLOCK i)))
Referenced model/specification definitions: Python source |
|---|
fused_layernorm_tritonfused_layernorm_triton_normalize_io_correctnessMathematical execution | Precomputed-input slice Normalization consumes supplied mean and inverse-variance scalars; their reduction is outside this IO slice. | Statement and preconditionsspecification fused_layernorm_triton_normalize_io_correctness
(out_ptr0 in_out_ptr0 in_ptr0 in_ptr1 in_ptr2 out_ptr1 : RegionName)
(rnumel RBLOCK : Nat) :
fusedNormalizeIO out_ptr0 in_out_ptr0 in_ptr0 in_ptr1 in_ptr2 out_ptr1 rnumel
RBLOCK
⊨ fun _p₀ _p₁ m1 m2 xs ws bs i => ((xs i - m1) * m2) * ws i + bs i
IO definition: fusedNormalizeIOdef fusedNormalizeIO
(out_ptr0 in_out_ptr0 in_ptr0 in_ptr1 in_ptr2 out_ptr1 : RegionName)
(rnumel RBLOCK : Nat) : Scalar2Tile3KernelIO where
kernel := fused_layernorm_triton_normalize_slice in_out_ptr0 in_ptr0 in_ptr1 in_ptr2 out_ptr0 out_ptr1 rnumel RBLOCK
sbuf1 := out_ptr0
sbuf2 := in_out_ptr0
tbuf1 := in_ptr0
tbuf2 := in_ptr1
tbuf3 := in_ptr2
out := out_ptr1
shape := [RBLOCK]
swin1 := fun p₀ _p₁ => p₀
swin2 := fun p₀ _p₁ => p₀
read1 := fun p₀ p₁ i => p₁ * RBLOCK + i.1.val + rnumel * p₀
read2 := fun _p₀ p₁ i => p₁ * RBLOCK + i.1.val
read3 := fun _p₀ p₁ i => p₁ * RBLOCK + i.1.val
write := fun p₀ p₁ i => p₁ * RBLOCK + i.1.val + rnumel * p₀
mask := fun _p₀ p₁ i => p₁ * RBLOCK + i.1.val < rnumel
Referenced model/specification definitions: Python source |
|---|
fused_recurrent_deltafused_recurrent_delta_output_summary_generalMathematical execution | Configured model / stage Forward/backward recurrence and step facts with state-carry premises. Conjuncts name distinct kernels and do not automatically yield a composed host IO contract. | Statement and preconditionsspecification fused_recurrent_delta_output_summary_general
(q k v beta o h0 ht dht dh0 do_ dq dk dv dbeta : RegionName)
(HPrev HOut DHPrev HRec : RegionName)
(IS_HEADWISE_BETA USE_INITIAL_STATE STORE_FINAL_STATE USE_DH0 USE_DHT : Bool)
(m s_qk_h s_vo_h NK B H T K V BK BV : Nat) (scale : ℝ) (s : BlockState)
(hBV : BV ≤ V) (hBK : BK ≤ K)
(hPrev : ∀ idx : TileIndex [BV, BK],
s.readMem HPrev (stateOffset s K V BK BV idx)
= deltaState s k v beta h0 IS_HEADWISE_BETA USE_INITIAL_STATE
s_qk_h s_vo_h T K V BK BV m idx)
-- the post-update carry buffer is the SAME region clause 4 stores into;
-- this hypothesis is NOT discharged by clause 4 (see the docstring)
(hNext : ∀ idx : TileIndex [BV, BK],
s.readMem HOut (stateOffset s K V BK BV idx)
= deltaState s k v beta h0 IS_HEADWISE_BETA USE_INITIAL_STATE
s_qk_h s_vo_h T K V BK BV (m + 1) idx) :
-- (1) the full forward surface lowers to the algorithm layer
(∃ alg, (fused_recurrent_delta_fwd_surface q k v beta o h0 ht s_qk_h s_vo_h
B H T K V BK BV scale USE_INITIAL_STATE STORE_FINAL_STATE
IS_HEADWISE_BETA).toAlgorithm? = Except.ok alg) ∧
-- (2) the full backward surface lowers to the algorithm layer
(∃ alg, (fused_recurrent_delta_bwd_surface q k v beta dht dh0 do_ dq dk dv
dbeta h0 s_qk_h s_vo_h NK B H T K V BK BV scale USE_INITIAL_STATE
IS_HEADWISE_BETA USE_DH0 USE_DHT).toAlgorithm? = Except.ok alg) ∧
-- (3) the `v_new` writeback realizes the genuine delta `vNewClosed(m)`
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_delta_vnew_step_slice HPrev k v
m s_qk_h s_vo_h K V BK BV)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jv : Fin BV => activeV s V BV jv)
(fun jv => (v, vRowOffset s m s_vo_h V BV jv)))
(expected := fun jv : Fin BV =>
vNewClosed s k v beta h0 IS_HEADWISE_BETA USE_INITIAL_STATE
s_qk_h s_vo_h T K V BK BV m jv)) ∧
-- (4) the state-update body realizes the genuine `deltaState(m+1)`
(ComputeCorrect.Realizes_without_Rounding
(kernel := if IS_HEADWISE_BETA then
fused_recurrent_delta_state_step_slice_headwise HPrev k v beta HOut
m s_qk_h s_vo_h T K V BK BV
else
fused_recurrent_delta_state_step_slice_scalarbeta HPrev k v beta HOut
m s_qk_h s_vo_h T K V BK BV)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BV, BK] => activeKV s K V BK BV idx)
(fun idx => (HOut, stateOffset s K V BK BV idx)))
(expected := fun idx =>
deltaState s k v beta h0 IS_HEADWISE_BETA USE_INITIAL_STATE
s_qk_h s_vo_h T K V BK BV (m + 1) idx)) ∧
-- (5) the output body realizes the genuine `outputClosed(m)` — mask-faithful,
-- so it holds for partial tiles and carries no full-tile antecedent
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_delta_output_step_slice HOut q o
m s_qk_h s_vo_h B H K V BK BV scale)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jv : Fin BV => activeV s V BV jv)
(fun jv => (o, outOffset s m s_vo_h B H V BV jv)))
(expected := fun jv : Fin BV =>
outputClosed s q k v beta h0 IS_HEADWISE_BETA USE_INITIAL_STATE
s_qk_h s_vo_h T K V BK BV scale m jv)) ∧
-- (6) the backward loop-1 `dk` body realizes the genuine `dkStepSpec` —
-- mask-faithful, so no full-tile antecedent
(ComputeCorrect.Realizes_without_Rounding
(kernel := if IS_HEADWISE_BETA then
fused_recurrent_delta_bwd_dk_step_slice_headwise DHPrev q do_ v beta dk
m s_qk_h s_vo_h B H T K V BK BV scale
else
fused_recurrent_delta_bwd_dk_step_slice_scalarbeta DHPrev q do_ v beta dk
m s_qk_h s_vo_h B H T K V BK BV scale)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jk : Fin BK => activeK s K BK jk)
(fun jk => (dk, dkRowOffset s m s_qk_h B H K BK jk)))
(expected := fun jk : Fin BK =>
dkStepSpec s DHPrev q do_ v beta IS_HEADWISE_BETA
m s_qk_h s_vo_h T K V BK BV scale jk)) ∧
-- (7) the backward loop-1 `dv` body realizes the genuine `dvStepSpec` —
-- mask-faithful, so no full-tile antecedent
(ComputeCorrect.Realizes_without_Rounding
(kernel := if IS_HEADWISE_BETA then
fused_recurrent_delta_bwd_dv_step_slice_headwise DHPrev q do_ k beta dv
m s_qk_h s_vo_h B H T K V BK BV scale
else
fused_recurrent_delta_bwd_dv_step_slice_scalarbeta DHPrev q do_ k beta dv
m s_qk_h s_vo_h B H T K V BK BV scale)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jv : Fin BV => activeV s V BV jv)
(fun jv => (dv, outOffset s m s_vo_h B H V BV jv)))
(expected := fun jv : Fin BV =>
dvStepSpec s DHPrev q do_ k beta IS_HEADWISE_BETA
m s_qk_h s_vo_h T K V BK BV scale jv)) ∧
-- (8) the backward loop-1 headwise `dbeta` body realizes `dbetaStepSpec` —
-- mask-faithful, so no full-tile antecedent
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_delta_bwd_dbeta_step_slice_headwise DHPrev q do_ k v
dbeta m s_qk_h s_vo_h NK B H T K V BK BV scale)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jv : Fin BV => activeV s V BV jv)
(fun jv => (dbeta, dbetaRowOffset s m s_vo_h B H NK V BV jv)))
(expected := fun jv : Fin BV =>
dbetaStepSpec s DHPrev q do_ k v m s_qk_h s_vo_h K V BK BV scale jv)) ∧
-- (9) the backward loop-1 scalar `dbeta` body realizes `dbetaScalarStepSpec`
-- — mask-faithful (guarded on both nested reductions)
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_delta_bwd_dbeta_step_slice_scalarbeta DHPrev q do_ k v
dbeta m s_qk_h s_vo_h B H T K V BK BV scale)
(initialState := s)
(write := fun _ : PUnit => some (dbeta, dbetaScalarOffset s m T B H))
(expected := fun _ =>
dbetaScalarStepSpec s DHPrev q do_ k v m s_qk_h s_vo_h K V BK BV scale)) ∧
-- (10) the backward loop-2 `dk` correction body, **sequenced onto the loop-1
-- stores**: given that the `dk`/`dv` rows still hold what clauses 6 and 7
-- put there, the final `dk` row is `dkCorrClosed`
((∀ jk : Fin BK,
s.readMem dk (dkRowOffset s m s_qk_h B H K BK jk)
= dkStepSpec s DHPrev q do_ v beta IS_HEADWISE_BETA
m s_qk_h s_vo_h T K V BK BV scale jk) →
(∀ jv : Fin BV,
s.readMem dv (outOffset s m s_vo_h B H V BV jv)
= dvStepSpec s DHPrev q do_ k beta IS_HEADWISE_BETA
m s_qk_h s_vo_h T K V BK BV scale jv) →
ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_delta_bwd_dk_correction_step_slice HRec dv dk
m s_qk_h s_vo_h B H K V BK BV)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jk : Fin BK => activeK s K BK jk)
(fun jk => (dk, dkRowOffset s m s_qk_h B H K BK jk)))
(expected := fun jk : Fin BK =>
dkCorrClosed s DHPrev HRec q do_ k v beta IS_HEADWISE_BETA
m s_qk_h s_vo_h T K V BK BV scale jk)) ∧
-- (11) the backward loop-2 `dq` body realizes the genuine `dqStepSpec` —
-- mask-faithful, so no full-tile antecedent
(ComputeCorrect.Realizes_without_Rounding
(kernel := if IS_HEADWISE_BETA then
fused_recurrent_delta_bwd_dq_step_slice_headwise HRec k v beta do_ dq
m s_qk_h s_vo_h B H T K V BK BV scale
else
fused_recurrent_delta_bwd_dq_step_slice_scalarbeta HRec k v beta do_ dq
m s_qk_h s_vo_h B H T K V BK BV scale)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jk : Fin BK => activeK s K BK jk)
(fun jk => (dq, dkRowOffset s m s_qk_h B H K BK jk)))
(expected := fun jk : Fin BK =>
dqStepSpec s HRec k v beta do_ IS_HEADWISE_BETA
m s_qk_h s_vo_h T K V BK BV scale jk))
Referenced model/specification definitions: Python source |
|---|
fused_recurrent_deltafused_recurrent_delta_output_step_io_correctnessMathematical execution | Precomputed-input slice Computes one output step from a supplied HNext state and q; the state update/preceding recurrence is outside this IO kernel. | Statement and preconditionsspecification fused_recurrent_delta_output_step_io_correctness
(HNext q o : RegionName) (t s_qk_h s_vo_h B H K V BK BV : Nat)
(hOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun jv : Fin BV =>
(p₂ + p₁ * B * H) * s_vo_h + p₀ * BV + jv.val + t * V)) :
outputStepIO HNext q o t s_qk_h s_vo_h B H K V BK BV
⊨ fun _p₀ p₁ xs ys jv =>
∑ jk : Fin BK,
if p₁ * BK + jk.val < K then
xs (jv.1, jk, PUnit.unit) * ys (jk, PUnit.unit)
else 0
IO definition: outputStepIOnoncomputable def outputStepIO (HNext q o : RegionName)
(t s_qk_h s_vo_h B H K V BK BV : Nat) :
Masked3DTileShapedKernelIO₂ where
kernel := fused_recurrent_delta_output_step_slice HNext q o t s_qk_h s_vo_h B H
K V BK BV 1
in1 := HNext
in2 := q
out := o
shape1 := [BV, BK]
shape2 := [BK]
shapeOut := [BV]
read1 := fun p₀ p₁ p₂ idx =>
p₂ * K * V + (p₁ * BK + idx.2.1.val) * V + (p₀ * BV + idx.1.val)
read2 := fun _p₀ p₁ p₂ jk => p₂ * s_qk_h + p₁ * BK + jk.1.val + t * K
write := fun p₀ p₁ p₂ jv =>
(p₂ + p₁ * B * H) * s_vo_h + p₀ * BV + jv.1.val + t * V
mask1 := fun p₀ p₁ _p₂ idx =>
p₁ * BK + idx.2.1.val < K ∧ p₀ * BV + idx.1.val < V
mask2 := fun _p₀ p₁ _p₂ jk => p₁ * BK + jk.1.val < K
writeMask := fun p₀ _p₁ _p₂ jv => p₀ * BV + jv.1.val < V
Referenced model/specification definitions: Python source |
|---|
fused_recurrent_hgrnfused_recurrent_hgrn_output_summary_generalMathematical execution | Configured model / stage Forward recurrence and backward step facts with explicit carried state. The selected backward IO headline is a single step. | Statement and preconditionsspecification fused_recurrent_hgrn_output_summary_general
(X G O H0 Ht DX DG DO BHPrev DHPrev : RegionName)
(USE_INITIAL_STATE STORE_FINAL_STATE : Bool)
(i_t T D BD : Nat) (s : BlockState) (hBD : 0 < BD)
(hPrev : ∀ i : Fin BD,
s.readMem BHPrev (bhOffset s D BD i)
= hgrnStateClosed s X G H0 USE_INITIAL_STATE T D BD i_t i) :
-- (1) the full forward surface lowers to the algorithm layer
(∃ alg, (fused_recurrent_hgrn_fwd_surface X G O H0 Ht T D BD
USE_INITIAL_STATE STORE_FINAL_STATE).toAlgorithm? = Except.ok alg) ∧
-- (2) the forward output body realizes the genuine `hgrnStateClosed(i_t+1)`
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_hgrn_forward_step_store_slice BHPrev X G O
i_t T D BD)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s D BD)
(fun i => (O, outOffset s i_t T D BD i)))
(expected := fun i =>
hgrnStateClosed s X G H0 USE_INITIAL_STATE T D BD (i_t + 1) i)) ∧
-- (3) the backward `dx` body realizes the genuine `dh_prev + do`
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_hgrn_bwd_dx_step_store_slice DHPrev DO DX
i_t T D BD)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s D BD)
(fun i => (DX, outOffset s i_t T D BD i)))
(expected := fun i => bwdDxStepValue s DHPrev DO i_t T D BD i)) ∧
-- (4) the backward `dg` body realizes the genuine `(dh_prev+do)·b_o`, with
-- Python's three-way previous-row `b_o` branch modeled
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_hgrn_bwd_dg_step_store_slice DHPrev DO O H0 DG
USE_INITIAL_STATE i_t T D BD)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s D BD)
(fun i => (DG, outOffset s i_t T D BD i)))
(expected := fun i =>
bwdDgStepValue s DHPrev DO O H0 USE_INITIAL_STATE i_t T D BD i)) ∧
-- (5) the backward scan step `b_dh = b_dh * b_g` executed: the two-step body
-- realizes the genuine closed form `hgrnBwdDx(i_t)` over `do`, `g`
(i_t + 1 < T →
(∀ i : Fin BD,
s.readMem DHPrev (outOffset s (i_t + 1) T D BD i)
= hgrnBwdCarry s DO G T D BD (i_t + 1) i) →
ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_hgrn_bwd_dx_two_step_store_slice DHPrev DO G DX
i_t T D BD)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s D BD)
(fun i => (DX, outOffset s i_t T D BD i)))
(expected := fun i => hgrnBwdDx s DO G T D BD i_t i)) ∧
-- (6) the same body at the loop's first fold, where the assumed carry is
-- only the `tl.zeros` seed
(2 ≤ T →
(∀ i : Fin BD, s.readMem DHPrev (outOffset s (T - 1) T D BD i) = 0) →
ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_hgrn_bwd_dx_two_step_store_slice DHPrev DO G DX
(T - 2) T D BD)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s D BD)
(fun i => (DX, outOffset s (T - 2) T D BD i)))
(expected := fun i => hgrnBwdDx s DO G T D BD (T - 2) i))
Referenced model/specification definitions: Python source |
|---|
fused_recurrent_hgrnfused_recurrent_hgrn_bwd_dx_io_correctnessMathematical execution | Precomputed-input slice Adds supplied DHPrev and DO at one backward timestep; production of the carried adjoint is outside this kernel. | Statement and preconditionsspecification fused_recurrent_hgrn_bwd_dx_io_correctness
(DHPrev DO DX : RegionName) (i_t T D BD : Nat) :
bwdDxIO DHPrev DO DX i_t T D BD
⊨ fun _p₀ _p₁ xs ys o => xs o + ys o
IO definition: bwdDxIOdef bwdDxIO (DHPrev DO DX : RegionName) (i_t T D BD : Nat) :
MaskedTileShapedKernelIO₂ where
kernel := fused_recurrent_hgrn_bwd_dx_step_store_slice DHPrev DO DX i_t T D BD
in1 := DHPrev
in2 := DO
out := DX
shape1 := [BD]
shape2 := [BD]
shapeOut := [BD]
read1 := fun p₀ p₁ i => (p₁ * T + i_t) * D + (p₀ * BD + i.1.val)
read2 := fun p₀ p₁ i => (p₁ * T + i_t) * D + (p₀ * BD + i.1.val)
write := fun p₀ p₁ o => (p₁ * T + i_t) * D + (p₀ * BD + o.1.val)
mask1 := fun p₀ _p₁ i => p₀ * BD + i.1.val < D
mask2 := fun p₀ _p₁ i => p₀ * BD + i.1.val < D
writeMask := fun p₀ _p₁ o => p₀ * BD + o.1.val < D
Referenced model/specification definitions: Python source |
|---|
fused_recurrent_retentionfused_recurrent_retention_output_summary_generalMathematical execution | Configured model / stage Forward/backward recurrence and step facts for the stated flags and tiles. Seed IO contracts cover initialization only. | Statement and preconditionsspecification fused_recurrent_retention_output_summary_general
(q k v o do_ dq dk dv initial_state final_state
HSeed HPrev HOut DHPrev DHOut : RegionName)
(USE_INITIAL_STATE STORE_FINAL_STATE : Bool)
(s_qk_h s_vo_h B H T DK DV BK BV m : Nat) (scale : ℝ) (s : BlockState)
(hBV : BV ≤ DV) (hBKpos : 0 < BK) (hBVpos : 0 < BV)
(hmT : m < T)
(hPrev : ∀ (jk : Fin BK) (jv : Fin BV),
s.readMem HPrev (stateOffset s DK DV BK BV jk jv)
= stateClosed s k v initial_state USE_INITIAL_STATE s_qk_h s_vo_h
H DK DV BK BV m jk jv)
(hDPrev : ∀ (jk : Fin BK) (jv : Fin BV),
s.readMem DHPrev (stateOffset s DK DV BK BV jk jv)
= bbVal s H *
dStateClosed s q do_ s_qk_h s_vo_h H DK DV BK BV T scale (m + 1) jk jv) :
-- (1) the full forward surface lowers to the algorithm layer
(∃ alg, (fused_recurrent_retention_fwd_surface q k v o initial_state
final_state s_qk_h s_vo_h B H T scale BK BV DK DV USE_INITIAL_STATE
STORE_FINAL_STATE).toAlgorithm? = Except.ok alg) ∧
-- (2) the full backward surface lowers to the algorithm layer
(∃ alg, (fused_recurrent_retention_bwd_surface q k v do_ dq dk dv
initial_state s_qk_h s_vo_h B H T scale BK BV DK DV
USE_INITIAL_STATE).toAlgorithm? = Except.ok alg) ∧
-- (2a) the seed prologue realizes `stateClosed 0` (the carry base case),
-- into the INTERNAL carry register `HSeed` (a fiction region)
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_retention_seed_slice initial_state HSeed
DK DV BK BV USE_INITIAL_STATE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BV, BK] => activeKV s DK DV BK BV idx)
(fun idx : TileIndex [BV, BK] =>
(HSeed, stateOffset s DK DV BK BV idx.2.1 idx.1)))
(expected := fun idx : TileIndex [BV, BK] =>
stateClosed s k v initial_state USE_INITIAL_STATE s_qk_h s_vo_h
H DK DV BK BV 0 idx.2.1 idx.1)) ∧
-- (2b) the reverse carry invariant's base case at the top of the scan
(∀ (jk : Fin BK) (jv : Fin BV),
bbVal s H *
dStateClosed s q do_ s_qk_h s_vo_h H DK DV BK BV T scale T jk jv = 0) ∧
-- (3) the forward output body realizes the genuine `outClosed(m)`
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_retention_output_step_slice HPrev q k v o
m s_qk_h s_vo_h B H DK DV BK BV scale)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jv : Fin BV => activeV s DV BV jv)
(fun jv => (o, outStepOffset s m s_vo_h B H DV BV jv)))
(expected := fun jv : Fin BV =>
outClosed s q k v initial_state USE_INITIAL_STATE s_qk_h s_vo_h
H DK DV BK BV scale m jv)) ∧
-- (4) the state-update body realizes the genuine `stateClosed(m+1)`
-- into the INTERNAL carry register `HOut` (a fiction region)
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_retention_state_step_slice HPrev k v HOut
m s_qk_h s_vo_h H DK DV BK BV)
(initialState := s)
(write := fun idx : TileIndex [BV, BK] =>
some (HOut, stateOffset s DK DV BK BV idx.2.1 idx.1))
(expected := fun idx =>
stateClosed s k v initial_state USE_INITIAL_STATE s_qk_h s_vo_h
H DK DV BK BV (m + 1) idx.2.1 idx.1)) ∧
-- (5) the backward `dq` body realizes the genuine `dqClosed(m)`
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_retention_bwd_dq_step_slice HPrev k v do_ dq
m s_qk_h s_vo_h B H DK DV BK BV scale)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jk : Fin BK => activeK s DK BK jk)
(fun jk => (dq, dqStepOffset s m s_qk_h B H DK BK jk)))
(expected := fun jk : Fin BK =>
dqClosed s k v do_ initial_state USE_INITIAL_STATE s_qk_h s_vo_h
H DK DV BK BV scale m jk)) ∧
-- (6) the reverse gradient-state carry body realizes `b_b·dStateClosed(m)`
-- into the INTERNAL carry register `DHOut` (a fiction region)
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_retention_bwd_dstate_step_slice DHPrev q do_
DHOut m s_qk_h s_vo_h H DK DV BK BV scale)
(initialState := s)
(write := fun idx : TileIndex [BK, BV] =>
some (DHOut, stateOffset s DK DV BK BV idx.1 idx.2.1))
(expected := fun idx =>
bbVal s H *
dStateClosed s q do_ s_qk_h s_vo_h H DK DV BK BV T scale m
idx.1 idx.2.1)) ∧
-- (7) the reverse `dk` body realizes the genuine `dkClosed(m)`
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_retention_bwd_dk_step_slice DHPrev q do_ v dk
m s_qk_h s_vo_h B H DK DV BK BV scale)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jk : Fin BK => activeK s DK BK jk)
(fun jk => (dk, dqStepOffset s m s_qk_h B H DK BK jk)))
(expected := fun jk : Fin BK =>
dkClosed s q do_ v s_qk_h s_vo_h H DK DV BK BV T scale m jk)) ∧
-- (8) the reverse `dv` body realizes the genuine `dvClosed(m)`
(ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_retention_bwd_dv_step_slice DHPrev q do_ k dv
m s_qk_h s_vo_h B H DK DV BK BV scale)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jv : Fin BV => activeV s DV BV jv)
(fun jv => (dv, outStepOffset s m s_vo_h B H DV BV jv)))
(expected := fun jv : Fin BV =>
dvClosed s q do_ k s_qk_h s_vo_h H DK DV BK BV T scale m jv))
Referenced model/specification definitions: Python source |
|---|
fused_recurrent_retentionfused_recurrent_retention_seed_io_correctnessMathematical execution | Configured model / stage Copies the supplied initial state into the seed tile with USE_INITIAL_STATE=true; this does not execute the recurrence. | Statement and preconditionsspecification fused_recurrent_retention_seed_io_correctness
(initial_state HSeed : RegionName) (DK DV BK BV : Nat)
(hOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BV, BK] =>
p₂ * DK * DV + (p₁ * BK + idx.2.1.val) * DV + (p₀ * BV + idx.1.val))) :
seedIO initial_state HSeed DK DV BK BV
⊨ fun _p₀ _p₁ xs idx => xs idx
IO definition: seedIOdef seedIO (initial_state HSeed : RegionName) (DK DV BK BV : Nat) :
Masked3DTileKernelIO₁ where
kernel := fused_recurrent_retention_seed_slice initial_state HSeed DK DV BK BV
Bool.true
inp := initial_state
out := HSeed
shape := [BV, BK]
read := fun p₀ p₁ p₂ idx =>
p₂ * DK * DV + (p₁ * BK + idx.2.1.val) * DV + (p₀ * BV + idx.1.val)
write := fun p₀ p₁ p₂ idx =>
p₂ * DK * DV + (p₁ * BK + idx.2.1.val) * DV + (p₀ * BV + idx.1.val)
mask := fun p₀ p₁ _p₂ idx =>
p₁ * BK + idx.2.1.val < DK ∧ p₀ * BV + idx.1.val < DV
Referenced model/specification definitions: Python source |
|---|
fused_recurrent_retentionfused_recurrent_retention_seed_io_correctnessRAbstract cast/store rounding | Configured model / stage Abstract rounding contract for the USE_INITIAL_STATE=true seed-copy stage only. | Statement and preconditionsspecification fused_recurrent_retention_seed_io_correctnessR (R : RoundingModel)
(initial_state HSeed : RegionName) (DK DV BK BV : Nat)
(hOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BV, BK] =>
p₂ * DK * DV + (p₁ * BK + idx.2.1.val) * DV + (p₀ * BV + idx.1.val))) :
seedIO initial_state HSeed DK DV BK BV
⊨[R, FloatDType.real] fun _p₀ _p₁ xs idx => xs idx
IO definition: seedIOdef seedIO (initial_state HSeed : RegionName) (DK DV BK BV : Nat) :
Masked3DTileKernelIO₁ where
kernel := fused_recurrent_retention_seed_slice initial_state HSeed DK DV BK BV
Bool.true
inp := initial_state
out := HSeed
shape := [BV, BK]
read := fun p₀ p₁ p₂ idx =>
p₂ * DK * DV + (p₁ * BK + idx.2.1.val) * DV + (p₀ * BV + idx.1.val)
write := fun p₀ p₁ p₂ idx =>
p₂ * DK * DV + (p₁ * BK + idx.2.1.val) * DV + (p₀ * BV + idx.1.val)
mask := fun p₀ p₁ _p₂ idx =>
p₁ * BK + idx.2.1.val < DK ∧ p₀ * BV + idx.1.val < DV
Referenced model/specification definitions: Python source |
|---|
fused_rotary_embeddingdecoding_fused_rotary_embedding_q_correctnessMathematical execution | Configured model / stage Separate Q-rotation and V-cache chain contracts. These selected sub-kernels do not by themselves prove the entire fused Q/K/V/cache program or arbitrary aliasing between stages. | Statement and preconditionsspecification decoding_fused_rotary_embedding_q_correctness
(q cos sin : RegionName)
(q_token_stride q_head_stride head_dim_stride cos_token_stride cos_stride
HALF_DIM : Nat)
(hstride : 0 < head_dim_stride) :
decodingRotaryQIO q cos sin q_token_stride q_head_stride head_dim_stride
cos_token_stride cos_stride HALF_DIM ⊨
fun _pid₀ _pid₁ xs o j =>
let q0 := xs (⟨0, by decide⟩ : Fin 4) j
let q1 := xs (⟨1, by decide⟩ : Fin 4) j
let c := xs (⟨2, by decide⟩ : Fin 4) j
let sn := xs (⟨3, by decide⟩ : Fin 4) j
match o with
| ⟨0, _⟩ => q0 * c - q1 * sn
| ⟨_ + 1, _⟩ => q0 * sn + q1 * c
IO definition: decodingRotaryQIOdef decodingRotaryQIO (q cos sin : RegionName)
(q_token_stride q_head_stride head_dim_stride cos_token_stride cos_stride
HALF_DIM : Nat) : GroupedMasked2DKernelIO where
kernel := decoding_fused_rotary_embedding_q_surface q cos sin q_token_stride
q_head_stride head_dim_stride cos_token_stride cos_stride (HALF_DIM * 2)
HALF_DIM
nIn := 4
nOut := 2
bufs := [q, cos, sin]
inp := fun
| ⟨0, _⟩ => q
| ⟨1, _⟩ => q
| ⟨2, _⟩ => cos
| ⟨_ + 3, _⟩ => sin
out := fun
| ⟨0, _⟩ => q
| ⟨_ + 1, _⟩ => q
B := HALF_DIM
read := fun
| ⟨0, _⟩ => fun p₀ p₁ j =>
qFirstAddr p₀ p₁ q_token_stride q_head_stride head_dim_stride HALF_DIM j
| ⟨1, _⟩ => fun p₀ p₁ j =>
qSecondAddr p₀ p₁ q_token_stride q_head_stride head_dim_stride HALF_DIM j
| ⟨2, _⟩ => fun _p₀ p₁ j =>
cosSinAddr p₁ cos_token_stride cos_stride HALF_DIM j
| ⟨_ + 3, _⟩ => fun _p₀ p₁ j =>
cosSinAddr p₁ cos_token_stride cos_stride HALF_DIM j
readMask := fun _ _ _ _ => True
write := fun
| ⟨0, _⟩ => fun p₀ p₁ j =>
qFirstAddr p₀ p₁ q_token_stride q_head_stride head_dim_stride HALF_DIM j
| ⟨_ + 1, _⟩ => fun p₀ p₁ j =>
qSecondAddr p₀ p₁ q_token_stride q_head_stride head_dim_stride HALF_DIM j
writeMask := fun _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
fused_rotary_embeddingdecoding_fused_rotary_embedding_vcache_chain_correctnessMathematical execution | Configured model / stage Separate Q-rotation and V-cache chain contracts. These selected sub-kernels do not by themselves prove the entire fused Q/K/V/cache program or arbitrary aliasing between stages. | Statement and preconditionsspecification decoding_fused_rotary_embedding_vcache_chain_correctness
(v v_cache : RegionName) (BLOCK_TABLES context_lengths : Region .nat)
(KV_GROUP_NUM k_token_stride k_head_stride head_dim_stride
vcb_stride vch_stride vcs_stride vcd_stride bts_stride btb_stride
block_size HEAD_DIM : Nat) :
decodingVCacheChainIO v v_cache BLOCK_TABLES context_lengths KV_GROUP_NUM
k_token_stride k_head_stride head_dim_stride vcb_stride vch_stride vcs_stride
vcd_stride bts_stride btb_stride block_size HEAD_DIM
⊨ decodingVCacheChainSpec
IO definition: decodingVCacheChainIOdef decodingVCacheChainIO
(v v_cache : RegionName) (BLOCK_TABLES context_lengths : Region .nat)
(KV_GROUP_NUM k_token_stride k_head_stride head_dim_stride
vcb_stride vch_stride vcs_stride vcd_stride bts_stride btb_stride
block_size HEAD_DIM : Nat) :
ChainMetaGroupedMasked2DKernelIO where
kernel := decoding_fused_rotary_embedding_vcache_chain v v_cache BLOCK_TABLES context_lengths KV_GROUP_NUM
k_token_stride k_head_stride head_dim_stride vcb_stride vch_stride vcs_stride
vcd_stride bts_stride btb_stride block_size HEAD_DIM
nIn := 1
nOut := 1
bufs := [v, v_cache, context_lengths, BLOCK_TABLES]
mbuf1 := context_lengths
mbuf2 := BLOCK_TABLES
inp := fun _ => v
out := fun _ => v_cache
B := HEAD_DIM
mwin1 := fun _ pid₁ => pid₁
mwin2 := fun _ pid₁ m₁ => pid₁ * bts_stride + ((m₁ - 1) / block_size) * btb_stride
read := fun _ pid₀ pid₁ _ _ j =>
pid₁ * k_token_stride + (pid₀ / KV_GROUP_NUM) * k_head_stride +
j.val * head_dim_stride
readMask := fun _ pid₀ _ _ _ _ => pid₀ % KV_GROUP_NUM = 0
write := fun _ pid₀ _ m₁ m₂ j =>
m₂ * vcb_stride + (pid₀ / KV_GROUP_NUM) * vch_stride +
((m₁ - 1) % block_size) * vcs_stride + j.val * vcd_stride
writeMask := fun _ pid₀ _ _ _ _ => pid₀ % KV_GROUP_NUM = 0
Referenced model/specification definitions: Python source |
|---|
fused_rwkv6_kernelfused_recurrent_rwkv6_output_summary_generalMathematical execution | Configured model / stage Forward recurrence, step, and conditional final-state facts. The state-step IO theorem starts from a supplied previous state. | Statement and preconditionsspecification fused_recurrent_rwkv6_output_summary_general
(q k v w u o h0 ht BHPrev BHOut C : RegionName)
(USE_INITIAL_STATE STORE_FINAL_STATE : Bool)
(s_k_h s_v_h B H T K V BK BV m : Nat) (scale : ℝ) (s sFinal : BlockState)
(hBV : BV ≤ V) (hBK : BK ≤ K) (hBVpos : 0 < BV) :
-- (1) the full surface lowers to the algorithm layer
(∃ alg, (fused_recurrent_rwkv6_fwd_surface q k v w u o h0 ht
s_k_h s_v_h B H T K V BK BV scale USE_INITIAL_STATE STORE_FINAL_STATE
Bool.false).toAlgorithm? = Except.ok alg) ∧
-- (2) the output body realizes the genuine `outputClosed(m)` — mask-faithful,
-- so this holds for partial key/value tiles as well. The pinned carry is a
-- clause-local antecedent, so it cannot weaken clause 4.
((∀ idx : TileIndex [BV, BK],
s.readMem BHPrev (finalStateOffset s K V BK BV idx)
= stateClosed s k v w h0 USE_INITIAL_STATE s_k_h s_v_h K V BK BV m idx) →
ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_rwkv6_output_step_slice BHPrev q k v u o
m s_k_h s_v_h B H T K V BK BV scale)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun jv : Fin BV => active s V BV jv)
(fun jv => (o, outStepOffset s m s_v_h B H V BV jv)))
(expected := fun jv : Fin BV =>
outputClosed s q k v w u h0 USE_INITIAL_STATE s_k_h s_v_h H K V BK BV
scale m jv)) ∧
-- (3) the state-update body realizes the genuine `stateClosed(m+1)`, under the
-- same clause-local pinned carry
((∀ idx : TileIndex [BV, BK],
s.readMem BHPrev (finalStateOffset s K V BK BV idx)
= stateClosed s k v w h0 USE_INITIAL_STATE s_k_h s_v_h K V BK BV m idx) →
ComputeCorrect.Realizes_without_Rounding
(kernel := fused_recurrent_rwkv6_state_step_slice BHPrev k v w BHOut
m s_k_h s_v_h K V BK BV)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BV, BK] => activeKV s K V BK BV idx)
(fun idx => (BHOut, finalStateOffset s K V BK BV idx)))
(expected := fun idx =>
stateClosed s k v w h0 USE_INITIAL_STATE s_k_h s_v_h K V BK BV
(m + 1) idx)) ∧
-- (4) the **cross-step carry fold**: chaining `T` state slices through one
-- shared carry region `C` reaches `stateClosed(T)` from a *single*
-- assumption about the initial buffer — no per-step pinned carry at all.
-- Its own antecedents keep clauses 2 and 3 free of them.
(k ≠ C → v ≠ C → w ≠ C →
(∀ idx : TileIndex [BV, BK], activeKV s K V BK BV idx →
s.readMem C (finalStateOffset s K V BK BV idx)
= stateSeed s h0 USE_INITIAL_STATE K V BK BV idx) →
execChain (foldStages
(fun j => fused_recurrent_rwkv6_state_step_slice C k v w C
j s_k_h s_v_h K V BK BV) T) s = some sFinal →
(AgreeOutsideRegion C s sFinal ∧
∀ idx : TileIndex [BV, BK], activeKV s K V BK BV idx →
sFinal.readMem C (finalStateOffset s K V BK BV idx)
= stateClosed s k v w h0 USE_INITIAL_STATE s_k_h s_v_h K V BK BV T idx))
Referenced model/specification definitions: Python source |
|---|
fused_rwkv6_kernelfused_recurrent_rwkv6_state_step_io_correctnessMathematical execution | Precomputed-input slice One state update from supplied BHPrev, k, v, and w; it does not construct BHPrev or compose all recurrent steps. | Statement and preconditionsspecification fused_recurrent_rwkv6_state_step_io_correctness
(BHPrev k v w BHOut : RegionName) (t s_k_h s_v_h K V BK BV : Nat)
(hOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BV, BK] =>
p₂ * K * V + (p₁ * BK + idx.2.1.val) * V + (p₀ * BV + idx.1.val))) :
stateStepIO BHPrev k v w BHOut t s_k_h s_v_h K V BK BV
⊨ fun _p₀ _p₁ xs ks vs ws idx =>
xs idx * Real.exp (ws (idx.2.1, PUnit.unit))
+ ks (idx.2.1, PUnit.unit) * vs (idx.1, PUnit.unit)
IO definition: stateStepIOnoncomputable def stateStepIO (BHPrev k v w BHOut : RegionName)
(t s_k_h s_v_h K V BK BV : Nat) : Masked3DTileShaped4KernelIO where
kernel := fused_recurrent_rwkv6_state_step_slice BHPrev k v w BHOut t s_k_h
s_v_h K V BK BV
in1 := BHPrev
in2 := k
in3 := v
in4 := w
out := BHOut
shape1 := [BV, BK]
shape2 := [BK]
shape3 := [BV]
shape4 := [BK]
shapeOut := [BV, BK]
read1 := fun p₀ p₁ p₂ idx =>
p₂ * K * V + (p₁ * BK + idx.2.1.val) * V + (p₀ * BV + idx.1.val)
read2 := fun _p₀ p₁ p₂ jk => p₂ * s_k_h + p₁ * BK + jk.1.val + t * K
read3 := fun p₀ _p₁ p₂ jv => p₂ * s_v_h + p₀ * BV + jv.1.val + t * V
read4 := fun _p₀ p₁ p₂ jk => p₂ * s_k_h + p₁ * BK + jk.1.val + t * K
write := fun p₀ p₁ p₂ idx =>
p₂ * K * V + (p₁ * BK + idx.2.1.val) * V + (p₀ * BV + idx.1.val)
mask1 := fun p₀ p₁ _p₂ idx =>
p₀ * BV + idx.1.val < V ∧ p₁ * BK + idx.2.1.val < K
mask2 := fun _p₀ p₁ _p₂ jk => p₁ * BK + jk.1.val < K
mask3 := fun p₀ _p₁ _p₂ jv => p₀ * BV + jv.1.val < V
mask4 := fun _p₀ p₁ _p₂ jk => p₁ * BK + jk.1.val < K
writeMask := fun p₀ p₁ _p₂ idx =>
p₀ * BV + idx.1.val < V ∧ p₁ * BK + idx.2.1.val < K
Referenced model/specification definitions: Python source |
|---|
geglu_tanh_tritongeglu_tanh_forward_kernel_correctnessMathematical execution | Configured model / stage Forward and backward tanh-GeGLU formulas under mathematical real semantics and the stated region-disjointness conditions; no concrete transcendental approximation claim. | Statement and preconditionsspecification geglu_tanh_forward_kernel_correctness
(A B C : RegionName)
(stride n_cols BLOCK_SIZE : Nat) :
gegluTanhFwdIO A B C stride n_cols BLOCK_SIZE
⊨ fun as bs i => TiledActivation.geluTanhFwd (as i) (bs i)
IO definition: gegluTanhFwdIOdef gegluTanhFwdIO (A B C : RegionName)
(stride n_cols BLOCK_SIZE : Nat) : MaskedKernelIO₂ where
kernel := geglu_tanh_forward_kernel A B C stride n_cols BLOCK_SIZE
in1 := A
in2 := B
out := C
B := BLOCK_SIZE
read1 := fun pid => pid * stride
read2 := fun pid => pid * stride
write := fun pid => pid * stride
mask := fun _ j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
geglu_tanh_tritongeglu_tanh_backward_kernel_correctnessMathematical execution | Configured model / stage Forward and backward tanh-GeGLU formulas under mathematical real semantics and the stated region-disjointness conditions; no concrete transcendental approximation claim. | Statement and preconditionsspecification geglu_tanh_backward_kernel_correctness
(DC A B : RegionName)
(stride n_cols BLOCK_SIZE : Nat)
(hAB : A ≠ B) :
gegluTanhBwdIO DC A B stride n_cols BLOCK_SIZE ⊨ fun dcs as bs =>
(fun i => TiledActivation.geluTanhBwdA (dcs i) (as i) (bs i),
fun i => TiledActivation.geluTanhBwdB (dcs i) (as i))
IO definition: gegluTanhBwdIOdef gegluTanhBwdIO (DC A B : RegionName)
(stride n_cols BLOCK_SIZE : Nat) : MaskedKernelIO₃ₓ₂ where
kernel := geglu_tanh_backward_kernel DC A B stride n_cols BLOCK_SIZE
bufs := [DC, A, B] -- a and b are updated in place
in1 := DC
in2 := A
in3 := B
out1 := A -- = in2: in-place `da` into `a`
out2 := B -- = in3: in-place `db` into `b`
B := BLOCK_SIZE
read1 := fun pid => pid * stride
read2 := fun pid => pid * stride
read3 := fun pid => pid * stride
write1 := fun pid => pid * stride
write2 := fun pid => pid * stride
mask := fun _ j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
index_select_bwdindex_select_cat_bwd_kernel_correctnessMathematical execution | Configured model / stage One indexed backward scatter of supplied gradients. Output-window validity/non-collision premises are required; duplicate-index concurrent accumulation is not implied. | Statement and preconditionsspecification index_select_cat_bwd_kernel_correctness
(grad_source_ptr index_ptr grad_output_ptr : RegionName)
(num_rows num_indices num_cols stride0 stride1
BLOCK_SIZE_INDEX BLOCK_SIZE_COL : Nat)
(hBC : 0 < BLOCK_SIZE_COL) :
indexSelectBwdIO grad_source_ptr index_ptr grad_output_ptr num_rows
num_indices num_cols stride0 stride1 BLOCK_SIZE_INDEX BLOCK_SIZE_COL ⊨
fun _ _ _ xs j => xs j
IO definition: indexSelectBwdIOdef indexSelectBwdIO (grad_source_ptr index_ptr grad_output_ptr : RegionName)
(num_rows num_indices num_cols stride0 stride1
BLOCK_SIZE_INDEX BLOCK_SIZE_COL : Nat) : GatherMasked2DKernelIO₁ where
kernel := index_select_cat_bwd_kernel grad_source_ptr index_ptr grad_output_ptr
num_rows num_indices num_cols stride0 stride1 BLOCK_SIZE_INDEX BLOCK_SIZE_COL
inp := grad_output_ptr
idxbuf := index_ptr
out := grad_source_ptr
B := BLOCK_SIZE_INDEX * BLOCK_SIZE_COL
readx := fun p₀ _ j => p₀ * BLOCK_SIZE_INDEX + j.val / BLOCK_SIZE_COL
read := fun p₀ p₁ _ j =>
(p₀ * BLOCK_SIZE_INDEX + j.val / BLOCK_SIZE_COL) * stride0 +
(p₁ * BLOCK_SIZE_COL + j.val % BLOCK_SIZE_COL) * stride1
write := fun _ p₁ ids j =>
ids j * stride0 + (p₁ * BLOCK_SIZE_COL + j.val % BLOCK_SIZE_COL) * stride1
mask := fun p₀ _ j =>
p₀ * BLOCK_SIZE_INDEX + j.val / BLOCK_SIZE_COL < num_indices
readMask := fun p₀ p₁ _ j =>
p₀ * BLOCK_SIZE_INDEX + j.val / BLOCK_SIZE_COL < num_indices ∧
p₁ * BLOCK_SIZE_COL + j.val % BLOCK_SIZE_COL < num_cols
writeMask := fun p₀ p₁ _ j =>
p₀ * BLOCK_SIZE_INDEX + j.val / BLOCK_SIZE_COL < num_indices ∧
p₁ * BLOCK_SIZE_COL + j.val % BLOCK_SIZE_COL < num_cols
Referenced model/specification definitions: Python source |
|---|
index_select_catindex_select_cat_fwd_kernel_correctnessMathematical execution | Configured model / stage One indexed gather/copy with natural-valued indices and positive column block size. IO bounds and layout assumptions apply; host concatenation is external. | Statement and preconditionsspecification index_select_cat_fwd_kernel_correctness
(output_ptr source_ptr index_ptr : RegionName)
(num_indices num_cols stride0 stride1
BLOCK_SIZE_INDEX BLOCK_SIZE_COL : Nat)
(hBC : 0 < BLOCK_SIZE_COL) :
indexSelectCatIO output_ptr source_ptr index_ptr num_indices num_cols
stride0 stride1 BLOCK_SIZE_INDEX BLOCK_SIZE_COL ⊨
fun _ _ _ xs j => xs j
IO definition: indexSelectCatIOdef indexSelectCatIO (output_ptr source_ptr index_ptr : RegionName)
(num_indices num_cols stride0 stride1
BLOCK_SIZE_INDEX BLOCK_SIZE_COL : Nat) : GatherMasked2DKernelIO₁ where
kernel := index_select_cat_fwd_kernel output_ptr source_ptr index_ptr
num_indices num_cols stride0 stride1 BLOCK_SIZE_INDEX BLOCK_SIZE_COL
inp := source_ptr
idxbuf := index_ptr
out := output_ptr
B := BLOCK_SIZE_INDEX * BLOCK_SIZE_COL
readx := fun p₀ _ j => p₀ * BLOCK_SIZE_INDEX + j.val / BLOCK_SIZE_COL
read := fun _ p₁ ids j =>
ids j * stride0 + (p₁ * BLOCK_SIZE_COL + j.val % BLOCK_SIZE_COL) * stride1
write := fun p₀ p₁ _ j =>
(p₀ * BLOCK_SIZE_INDEX + j.val / BLOCK_SIZE_COL) * stride0 +
(p₁ * BLOCK_SIZE_COL + j.val % BLOCK_SIZE_COL) * stride1
mask := fun p₀ _ j =>
p₀ * BLOCK_SIZE_INDEX + j.val / BLOCK_SIZE_COL < num_indices
readMask := fun p₀ p₁ _ j =>
p₀ * BLOCK_SIZE_INDEX + j.val / BLOCK_SIZE_COL < num_indices ∧
p₁ * BLOCK_SIZE_COL + j.val % BLOCK_SIZE_COL < num_cols
writeMask := fun p₀ p₁ _ j =>
p₀ * BLOCK_SIZE_INDEX + j.val / BLOCK_SIZE_COL < num_indices ∧
p₁ * BLOCK_SIZE_COL + j.val % BLOCK_SIZE_COL < num_cols
Referenced model/specification definitions: Python source |
|---|
int4_matmulint4_matmul_exec_genuineMathematical execution | Configured model / stage Selected packed-int4 mathematical decode/matmul surface with the displayed grouping/block assumptions. The exec-existential omits a frame and does not certify fixed-width overflow behavior. | Statement and preconditionsspecification int4_matmul_exec_genuine
(a_ptr c_ptr bs_ptr : RegionName) (b_ptr bzp_ptr : Region .nat)
(M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
stride_bsk stride_bsn stride_bzpk stride_bzpn group_size : Nat)
(BM BN BK GM numKBlocks : Nat) (s : BlockState)
(hK : K = BK * numKBlocks)
(hBK8 : BK % 8 = 0)
(hpid1 : s.pids 1 = 0)
(hInj : Function.Injective
(fun i : TileIndex [BM, BN] => cAddr stride_cm stride_cn BM BN
(pidM s M N BM BN GM) (pidN s M N BM BN GM) i)) :
∃ sF, exec (int4_matmul_surface a_ptr c_ptr bs_ptr b_ptr bzp_ptr
M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
stride_bsk stride_bsn stride_bzpk stride_bzpn group_size
BM BN BK GM numKBlocks).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BM, BN],
(pidM s M N BM BN GM * BM + idx.1.val < M
∧ pidN s M N BM BN GM * BN + idx.2.1.val < N) →
sF.readMem c_ptr (cAddr stride_cm stride_cn BM BN (pidM s M N BM BN GM)
(pidN s M N BM BN GM) idx)
= accSpec s a_ptr bs_ptr b_ptr bzp_ptr group_size stride_am
stride_ak stride_bk stride_bn stride_bsk stride_bsn
stride_bzpk stride_bzpn BK numKBlocks
(pidM s M N BM BN GM * BM + idx.1.val)
(pidN s M N BM BN GM * BN + idx.2.1.val)
Referenced model/specification definitions: Python source |
|---|
int8_dequant_matmulint8_dequant_matmul_exec_genuineMathematical execution | Configured model / stage Selected rowwise dequantizing matmul with optional bias and a typed fp16 output. The exec-existential omits a frame; its mathematical casts are not hardware IEEE conversion. | Statement and preconditionsspecification int8_dequant_matmul_exec_genuine
(A B : Region .int) (C bias state_x_ptr state_w_ptr : RegionName)
(M N K : Nat) (divfactor : ℝ) (has_bias : Bool)
(stride_am stride_ak stride_bk stride_bn stride_cm stride_cn : Nat)
(BM BN BK GM numKBlocks : Nat) (s : BlockState)
(hK : K = BK * numKBlocks)
(hpid1 : s.pids 1 = 0)
(hInj : Function.Injective
(fun i : TileIndex [BM, BN] => cAddr stride_cm stride_cn BM BN
(pidM s M N BM BN GM) (pidN s M N BM BN GM) i)) :
∃ sF, exec (int8_dequant_matmul_surface A B C bias state_x_ptr state_w_ptr
M N K divfactor has_bias
stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
BM BN BK GM numKBlocks).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BM, BN],
(pidM s M N BM BN GM * BM + idx.1.val < M
∧ pidN s M N BM BN GM * BN + idx.2.1.val < N) →
sF.mem C (cAddr stride_cm stride_cn BM BN (pidM s M N BM BN GM)
(pidN s M N BM BN GM) idx)
= MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (i8Spec s A B bias state_x_ptr state_w_ptr
stride_am stride_ak stride_bk stride_bn BK numKBlocks
divfactor has_bias
(pidM s M N BM BN GM * BM + idx.1.val)
(pidN s M N BM BN GM * BN + idx.2.1.val))))
Referenced model/specification definitions: Python source |
|---|
int8_matmul_kernelint8_matmul_kernel_exec_genuineMathematical execution | Configured model / stage Integer matmul over the modeled integer carrier and configured tile/reduction layout. The exec-existential omits a frame and does not establish machine overflow semantics. | Statement and preconditionsspecification int8_matmul_kernel_exec_genuine
(A : Region .int) (B : Region .nat) (C : RegionName)
(M N K : Nat)
(stride_am stride_ak stride_bk stride_bn stride_cm stride_cn : Nat)
(BM BN BK GM numKBlocks : Nat) (s : BlockState)
(hK : K = 4 * (BK * numKBlocks))
(hInj : Function.Injective
(fun i : TileIndex [BM, BN] => imCAddr stride_cm stride_cn BM BN
(pidM s M N BM BN GM) (pidN s M N BM BN GM) i)) :
∃ sF, exec (int8_matmul_kernel_surface A B C M N K
stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
BM BN BK GM numKBlocks).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BM, BN],
(pidM s M N BM BN GM * BM + idx.1.val < M
∧ pidN s M N BM BN GM * BN + idx.2.1.val < N) →
sF.mem C (imCAddr stride_cm stride_cn BM BN (pidM s M N BM BN GM)
(pidN s M N BM BN GM) idx)
= MemCell.of .int (imSpec s A B stride_am stride_ak stride_bk
stride_bn BK numKBlocks
(pidM s M N BM BN GM * BM + idx.1.val)
(pidN s M N BM BN GM * BN + idx.2.1.val))
Referenced model/specification definitions: Python source |
|---|
int8_matmul_quantizationint8_matmul_quantization_quantize_exec_genuineMathematical execution | Configured model / stage Separate per-row quantization and matmul stages using the defined mathematical integer/cast operations. No frame conjunct or proof of end-to-end hardware quantization/overflow behavior. | Statement and preconditionsspecification int8_matmul_quantization_quantize_exec_genuine
(fpa_ptr : RegionName) (a_ptr : Region .int) (as_ptr : RegionName)
(M K : Nat)
(stride_fpam stride_fpak stride_am stride_ak stride_asm : Nat)
(BM BK numKBlocks : Nat) (s : BlockState)
(hBK : 0 < BK)
(hK : K ≤ numKBlocks * BK)
(hFit : s.pids 0 * BM + BM ≤ M)
(hFpaNe : fpa_ptr ≠ (Region.cast a_ptr : RegionName))
(hAsNe : as_ptr ≠ (Region.cast a_ptr : RegionName))
(hInj : Function.Injective (fun p : Fin BM × Fin K =>
(s.pids 0 * BM + p.1.val) * stride_am + p.2.val * stride_ak)) :
∃ sF, exec (int8_matmul_quantization_quantize_surface fpa_ptr a_ptr as_ptr
M K stride_fpam stride_fpak stride_am stride_ak stride_asm
BM BK numKBlocks).toAlgKernel s = some sF
∧ (∀ (r : Fin BM) (kg : Fin K),
sF.mem (Region.cast a_ptr)
((s.pids 0 * BM + r.val) * stride_am + kg.val * stride_ak)
= MemCell.of .int (qInt8Spec K BK numKBlocks
(qFpaElem s fpa_ptr stride_fpam stride_fpak
(s.pids 0 * BM + r.val)) kg.val))
∧ (∀ r : Fin BM,
sF.mem as_ptr (s.pids 0 * BM * stride_asm + r.val)
= MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (qScaleSpec K BK numKBlocks
(qFpaElem s fpa_ptr stride_fpam stride_fpak
(s.pids 0 * BM + r.val))))))
Referenced model/specification definitions: Python source |
|---|
int8_matmul_quantizationint8_matmul_quantization_matmul_exec_genuineMathematical execution | Configured model / stage Separate per-row quantization and matmul stages using the defined mathematical integer/cast operations. No frame conjunct or proof of end-to-end hardware quantization/overflow behavior. | Statement and preconditionsspecification int8_matmul_quantization_matmul_exec_genuine
(A : Region .int) (as_ptr : RegionName) (B : Region .int) (bs_ptr C : RegionName)
(M N K : Nat)
(stride_am stride_ak stride_asm stride_bk stride_bn stride_bsn
stride_cm stride_cn : Nat)
(BM BN BK GM numKBlocks : Nat) (s : BlockState)
(hK : K ≤ numKBlocks * BK)
(hpid1 : s.pids 1 = 0)
(hInj : Function.Injective
(fun i : TileIndex [BM, BN] => mmCAddr stride_cm stride_cn BM BN
(mmPidM s M N BM BN GM) (mmPidN s M N BM BN GM) i)) :
∃ sF, exec (int8_matmul_quantization_matmul_surface A as_ptr B bs_ptr C
M N K stride_am stride_ak stride_asm stride_bk stride_bn stride_bsn
stride_cm stride_cn BM BN BK GM numKBlocks).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BM, BN],
(mmPidM s M N BM BN GM * BM + idx.1.val < M
∧ mmPidN s M N BM BN GM * BN + idx.2.1.val < N) →
sF.mem C (mmCAddr stride_cm stride_cn BM BN (mmPidM s M N BM BN GM)
(mmPidN s M N BM BN GM) idx)
= MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some ((((∑ j : Fin K, mmAElem s A stride_am stride_ak
(mmPidM s M N BM BN GM * BM + idx.1.val) j.val
* mmBElem s B stride_bk stride_bn j.val
(mmPidN s M N BM BN GM * BN + idx.2.1.val)) : ℤ) : ℝ)
* mmAScaleElem s as_ptr stride_asm
(mmPidM s M N BM BN GM * BM + idx.1.val)
* mmBScaleElem s bs_ptr stride_bsn
(mmPidN s M N BM BN GM * BN + idx.2.1.val))))
Referenced model/specification definitions: Python source |
|---|
int8_quantizationper_block_int8_correctnessMathematical execution | Precomputed-input slice Q/K store slice with ScalePre input. Scale reduction, sign-dependent bias, and final int8 rounding are not part of the value contract. | Statement and preconditionsspecification per_block_int8_correctness
(X ScalePre XInt8 Scale : RegionName)
(L C BLK scale_stride : Nat) (preScale : ℝ)
(hB : 0 < BLK * C) (hRegions : XInt8 ≠ Scale) :
perBlockInt8IO X ScalePre XInt8 Scale L C BLK scale_stride preScale ⊨
fun _ _ xs ys =>
(fun j => perBlockInt8ValSpec BLK C preScale xs ys j, fun j => ys j)
IO definition: perBlockInt8IOdef perBlockInt8IO (X ScalePre XInt8 Scale : RegionName)
(L C BLK scale_stride : Nat) (preScale : ℝ) : Masked2DKernelIO₂ₓ₂ where
kernel := per_block_int8_store_slice X ScalePre XInt8 Scale L C BLK
scale_stride preScale
in1 := X
in2 := ScalePre
out1 := XInt8
out2 := Scale
B := BLK * C
read1 := fun off_blk off_b j => off_b * L * C + off_blk * BLK * C + j.val
read2 := fun off_blk off_b _ => off_b * scale_stride + off_blk
write1 := fun off_blk off_b j => off_b * L * C + off_blk * BLK * C + j.val
write2 := fun off_blk off_b _ => off_b * scale_stride + off_blk
mask := fun off_blk _ j => off_blk * BLK + j.val / C < L
read2Mask := fun _ _ _ => True
writeMask2 := fun _ _ j => j.val = 0
Referenced model/specification definitions: Python source |
|---|
int_scaled_matmulint_scaled_matmul_matmul_exec_genuineMathematical execution | Configured model / stage Separate integer matmul and scaled-output stages for explicit block-pointer layouts. Exec-existential readbacks omit a frame; integer arithmetic is not a fixed-width overflow certificate. | Statement and preconditionsspecification int_scaled_matmul_matmul_exec_genuine
(a_ptr b_ptr : Region .int) (c_ptr : RegionName)
(M N K : Nat)
(stride_am stride_ak stride_bk stride_bn stride_cm stride_cn : Nat)
(BM BN BK GM : Nat) (s : BlockState)
(hBK : 0 < BK)
(hInj : Function.Injective
(fun i : TileIndex [BM, BN] => mmCAddr stride_cm stride_cn BM BN
(mmPidM s M N BM BN GM) (mmPidN s M N BM BN GM) i)) :
∃ sF, exec (int_scaled_matmul_matmul_surface a_ptr b_ptr c_ptr M N K
stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
BM BN BK GM).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BM, BN],
(mmPidM s M N BM BN GM * BM + idx.1.val < M
∧ mmPidN s M N BM BN GM * BN + idx.2.1.val < N) →
sF.mem c_ptr (mmCAddr stride_cm stride_cn BM BN (mmPidM s M N BM BN GM)
(mmPidN s M N BM BN GM) idx)
= MemCell.of .int (mmSpec s a_ptr b_ptr K stride_am stride_ak
stride_bk stride_bn
(mmPidM s M N BM BN GM * BM + idx.1.val)
(mmPidN s M N BM BN GM * BN + idx.2.1.val))
Referenced model/specification definitions: Python source |
|---|
int_scaled_matmulint_scaled_matmul_scaled_exec_genuineMathematical execution | Configured model / stage Separate integer matmul and scaled-output stages for explicit block-pointer layouts. Exec-existential readbacks omit a frame; integer arithmetic is not a fixed-width overflow certificate. | Statement and preconditionsspecification int_scaled_matmul_scaled_exec_genuine
(a_ptr b_ptr : Region .int) (c_ptr s1_ptr : RegionName)
(M N K : Nat)
(stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
stride_s1m stride_s1n : Nat)
(BM BN BK GM : Nat) (EVEN_K : Bool) (numKBlocks : Nat) (s : BlockState)
(hCeil : K ≤ numKBlocks * BK)
(hEven : EVEN_K = Bool.true → K = numKBlocks * BK) :
∃ sF, exec (int_scaled_matmul_scaled_surface a_ptr b_ptr c_ptr s1_ptr M N K
stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
stride_s1m stride_s1n BM BN BK GM EVEN_K numKBlocks).toAlgKernel s
= some sF
∧ ∀ idx : TileIndex [BM, BN],
(smPidM s M N BM BN GM * BM + idx.1.val < M
∧ smPidN s M N BM BN GM * BN + idx.2.1.val < N) →
sF.mem c_ptr (smCAddr N BM BN (smPidM s M N BM BN GM)
(smPidN s M N BM BN GM) idx)
= MemCell.of .real (some
(((smSpec s a_ptr b_ptr K stride_am stride_ak stride_bk
stride_bn
(smPidM s M N BM BN GM * BM + idx.1.val)
(smPidN s M N BM BN GM * BN + idx.2.1.val) : ℤ) : ℝ)
* smS1 s s1_ptr (smPidM s M N BM BN GM * BM + idx.1.val)))
Referenced model/specification definitions: Python source |
|---|
iv_dependent_matmuliv_dependent_matmul_closed_form_correctMathematical execution | Configured model / stage Pre-load matmul surface with K=BK*numKBlocks and the stated contiguous/injective output layout. Other loop/launch configurations are not inferred. | Statement and preconditionsspecification iv_dependent_matmul_closed_form_correct
(A B C : RegionName) (s : BlockState)
(M N SAM SAK SBK SBN SCM SCN BM BN BK numKBlocks : Nat) (hBK : 0 < BK)
(hcn : SCN = 1) (hbnle : BN ≤ SCM)
(hmlt : ∀ i : Fin BM, rowIndex (pidM (s.pids 0) N BN) BM i < M)
(hnlt : ∀ j : Fin BN, colIndex (pidN (s.pids 0) N BN) BN j < N)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := iv_dependent_matmul_pre_load_surface A B C M N (BK * numKBlocks) SAM SAK SBK SBN SCM SCN BM BN BK)
(initialState := s)
(write := fun idx : TileIndex [BM, BN] =>
some (C, cOffset s (pidM (s.pids 0) N BN) (pidN (s.pids 0) N BN) BM BN SCM SCN idx))
(expected := fun idx : TileIndex [BM, BN] =>
outputCell s A B (pidM (s.pids 0) N BN) (pidN (s.pids 0) N BN)
BM BN M N SAM SAK SBK SBN BK numKBlocks idx)
Referenced model/specification definitions: Python source |
|---|
iv_dependent_matmuliv_dependent_matmul_io_correctnessAbstract cast/store rounding | Configured model / stage Pre-load matmul surface with K=BK*numKBlocks and the stated contiguous/injective output layout. Other loop/launch configurations are not inferred. | Statement and preconditionsspecification iv_dependent_matmul_io_correctness (R : RoundingModel)
(A B C : RegionName)
(M N SAM SAK SBK SBN SCM SCN BM BN BK numKBlocks : Nat)
(hBK : 0 < BK) (hcn : SCN = 1) (hbnle : BN ≤ SCM) :
ivDependentMatmulIO A B C M N SAM SAK SBK SBN SCM SCN BM BN BK numKBlocks
⊨[R] fun _ _ xs ys l =>
∑ t : Fin numKBlocks, ∑ e : Fin BK,
xs t (aLane BM BN BK l e) * ys t (bLane BM BN BK l e)
IO definition: ivDependentMatmulIOdef ivDependentMatmulIO (A B C : RegionName)
(M N SAM SAK SBK SBN SCM SCN BM BN BK numKBlocks : Nat) :
StreamMasked2DKernelIO₂ where
kernel := iv_dependent_matmul_pre_load_surface A B C M N (BK * numKBlocks) SAM SAK SBK SBN
SCM SCN BM BN BK
inp1 := A
inp2 := B
out := C
T := numKBlocks
B1 := BM * BK
B2 := BK * BN
C := BM * BN
outDType := .fp16
read1 := fun p₀ _ t l =>
(pidM p₀ N BN * BM + l.val / BK) % M * SAM + (t.val * BK + l.val % BK) * SAK
read2 := fun p₀ _ t l =>
(t.val * BK + l.val / BN) * SBK + (pidN p₀ N BN * BN + l.val % BN) % N * SBN
write := fun p₀ _ l =>
SCM * (pidM p₀ N BN * BM + l.val / BN) + SCN * (pidN p₀ N BN * BN + l.val % BN)
mask1 := fun _ _ t l => t.val * BK + l.val % BK < BK * numKBlocks
mask2 := fun _ _ t l => t.val * BK + l.val / BN < BK * numKBlocks
writeMask := fun p₀ _ l =>
pidM p₀ N BN * BM + l.val / BN < M ∧ pidN p₀ N BN * BN + l.val % BN < N
Referenced model/specification definitions: Python source |
|---|
kcache_copy_tritoncopy_to_kcache_seqlen_n1_correctnessMathematical execution | Configured model / stage Sequence-length-one cache-copy specialization with the explicit block-table and cache layout. It does not cover an arbitrary multi-token Python launch. | Statement and preconditionsspecification copy_to_kcache_seqlen_n1_correctness
(K KCache BLOCK_TABLES seq_lengths : RegionName)
(SPLIT_X stride_kt stride_kh stride_kd stride_kcb stride_kch
stride_kcsplit_x stride_kcs stride_bts stride_btb block_size
KCACHE_X : Nat) :
kcacheCopyN1IO K KCache BLOCK_TABLES seq_lengths SPLIT_X stride_kt
stride_kh stride_kd stride_kcb stride_kch stride_kcsplit_x stride_kcs
stride_bts stride_btb block_size KCACHE_X ⊨
fun _ _ _ _ xs _ => (xs, xs)
IO definition: kcacheCopyN1IOdef kcacheCopyN1IO (K KCache BLOCK_TABLES seq_lengths : RegionName)
(SPLIT_X stride_kt stride_kh stride_kd stride_kcb stride_kch
stride_kcsplit_x stride_kcs stride_bts stride_btb block_size
KCACHE_X : Nat) : ChainMetaMasked2DKernelIO₂ₓ₂ where
kernel := copy_to_kcache_seqlen_n1_surface K KCache BLOCK_TABLES seq_lengths
SPLIT_X stride_kt stride_kh stride_kd stride_kcb stride_kch
stride_kcsplit_x stride_kcs 0 stride_bts stride_btb block_size 1 0 KCACHE_X
mbuf1 := seq_lengths
mbuf2 := BLOCK_TABLES
in1 := K
in2 := K
out1 := KCache
out2 := KCache
B := KCACHE_X
mwin1 := fun pid₀ _ => pid₀
mwin2 := fun pid₀ _ m₁ =>
pid₀ * stride_bts + ((m₁ - 1) / block_size) * stride_btb
read1 := fun pid₀ pid₁ _ _ j =>
pid₀ * stride_kt + pid₁ * stride_kh +
(SPLIT_X * KCACHE_X + j.val) * stride_kd
read2 := fun pid₀ pid₁ _ _ j =>
pid₀ * stride_kt + pid₁ * stride_kh +
(SPLIT_X * KCACHE_X + j.val) * stride_kd
mask1 := fun _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ => False
write1 := fun _ pid₁ m₁ m₂ j =>
m₂ * stride_kcb + pid₁ * stride_kch + SPLIT_X * stride_kcsplit_x +
((m₁ - 1) % block_size) * stride_kcs + j.val
write2 := fun _ pid₁ m₁ m₂ j =>
m₂ * stride_kcb + pid₁ * stride_kch + SPLIT_X * stride_kcsplit_x +
((m₁ - 1) % block_size) * stride_kcs + j.val
writeMask1 := fun _ _ _ _ _ => True
writeMask2 := fun _ _ _ _ _ => False
Referenced model/specification definitions: Python source |
|---|
kldiv_computekldivergence_kernel_correctnessMathematical execution | Configured model / stage Per-program KL expression on mathematical inputs. Log/division use the model definitions; domain assumptions and hardware exceptional values require separate interpretation. | Statement and preconditionsspecification kldivergence_kernel_correctness
(x_ptr y_ptr output_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) :
kldivIO x_ptr y_ptr output_ptr n_elements BLOCK_SIZE
⊨ fun xs ys i => klDivSpec (xs i) (ys i)
IO definition: kldivIOdef kldivIO (x_ptr y_ptr output_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) : MaskedKernelIO₂ where
kernel := kldivergence_kernel x_ptr y_ptr output_ptr n_elements BLOCK_SIZE
in1 := x_ptr
in2 := y_ptr
out := output_ptr
B := BLOCK_SIZE
read1 := fun pid => pid * BLOCK_SIZE
read2 := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < n_elements
Referenced model/specification definitions: Python source |
|---|
kldiv_opskldiv_backward_default_compute_correctMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. Other reduction modes and host gradient scaling are not implied by these branch contracts. | Statement and preconditionsspecification kldiv_backward_default_compute_correct
(target_ptr new_grads_ptr : RegionName)
(target_stride new_grads_stride n_cols BLOCK_SIZE : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => linearOffset s new_grads_stride i)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := kldiv_backward_default target_ptr new_grads_ptr
target_stride new_grads_stride n_cols BLOCK_SIZE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < n_cols)
(fun i => (new_grads_ptr, linearOffset s new_grads_stride i)))
(expected := fun i => defaultSpec s target_ptr target_stride i)
Referenced model/specification definitions: Python source |
|---|
kldiv_opskldiv_backward_log_target_compute_correctMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. Other reduction modes and host gradient scaling are not implied by these branch contracts. | Statement and preconditionsspecification kldiv_backward_log_target_compute_correct
(target_ptr new_grads_ptr : RegionName)
(target_stride new_grads_stride n_cols BLOCK_SIZE : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => linearOffset s new_grads_stride i)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := kldiv_backward_log_target target_ptr new_grads_ptr
target_stride new_grads_stride n_cols BLOCK_SIZE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < n_cols)
(fun i => (new_grads_ptr, linearOffset s new_grads_stride i)))
(expected := fun i => logTargetSpec s target_ptr target_stride i)
Referenced model/specification definitions: Python source |
|---|
kldiv_opskldiv_forward_default_none_compute_correctMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. Other reduction modes and host gradient scaling are not implied by these branch contracts. | Statement and preconditionsspecification kldiv_forward_default_none_compute_correct
(y_ptr gt_ptr loss_ptr : RegionName)
(y_stride gt_stride loss_stride n_cols BLOCK_SIZE : Nat) (eps : ℝ)
(s : BlockState)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => linearOffset s loss_stride i)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := kldiv_forward_default_none y_ptr gt_ptr loss_ptr
y_stride gt_stride loss_stride n_cols BLOCK_SIZE eps)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < n_cols)
(fun i => (loss_ptr, linearOffset s loss_stride i)))
(expected := fun i =>
forwardDefaultSpec s y_ptr gt_ptr y_stride gt_stride eps i)
Referenced model/specification definitions: Python source |
|---|
kldiv_opskldiv_forward_log_target_none_compute_correctMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. Other reduction modes and host gradient scaling are not implied by these branch contracts. | Statement and preconditionsspecification kldiv_forward_log_target_none_compute_correct
(y_ptr gt_ptr loss_ptr : RegionName)
(y_stride gt_stride loss_stride n_cols BLOCK_SIZE : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => linearOffset s loss_stride i)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := kldiv_forward_log_target_none y_ptr gt_ptr loss_ptr
y_stride gt_stride loss_stride n_cols BLOCK_SIZE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < n_cols)
(fun i => (loss_ptr, linearOffset s loss_stride i)))
(expected := fun i =>
forwardLogTargetSpec s y_ptr gt_ptr y_stride gt_stride i)
Referenced model/specification definitions: Python source |
|---|
kldiv_opskldiv_forward_log_target_none_correctnessMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. Other reduction modes and host gradient scaling are not implied by these branch contracts. | Statement and preconditionsspecification kldiv_forward_log_target_none_correctness
(y_ptr gt_ptr loss_ptr : RegionName)
(y_stride gt_stride loss_stride n_cols BLOCK_SIZE : Nat) :
kldivForwardLogTargetIO y_ptr gt_ptr loss_ptr
y_stride gt_stride loss_stride n_cols BLOCK_SIZE
⊨ fun ys gts i => Real.exp (gts i) * (gts i - ys i)
IO definition: kldivForwardLogTargetIOdef kldivForwardLogTargetIO (y_ptr gt_ptr loss_ptr : RegionName)
(y_stride gt_stride loss_stride n_cols BLOCK_SIZE : Nat) : MaskedKernelIO₂ where
kernel := kldiv_forward_log_target_none y_ptr gt_ptr loss_ptr
y_stride gt_stride loss_stride n_cols BLOCK_SIZE
in1 := y_ptr
in2 := gt_ptr
out := loss_ptr
B := BLOCK_SIZE
read1 := fun pid => pid * y_stride
read2 := fun pid => pid * gt_stride
write := fun pid => pid * loss_stride
mask := fun _pid j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
kldiv_opskldiv_forward_default_none_correctnessMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. Other reduction modes and host gradient scaling are not implied by these branch contracts. | Statement and preconditionsspecification kldiv_forward_default_none_correctness
(y_ptr gt_ptr loss_ptr : RegionName)
(y_stride gt_stride loss_stride n_cols BLOCK_SIZE : Nat) (eps : ℝ) :
kldivForwardDefaultIO y_ptr gt_ptr loss_ptr
y_stride gt_stride loss_stride n_cols BLOCK_SIZE eps
⊨ fun ys gts i => gts i * (Real.log (max (gts i) eps) - ys i)
IO definition: kldivForwardDefaultIOdef kldivForwardDefaultIO (y_ptr gt_ptr loss_ptr : RegionName)
(y_stride gt_stride loss_stride n_cols BLOCK_SIZE : Nat) (eps : ℝ) : MaskedKernelIO₂ where
kernel := kldiv_forward_default_none y_ptr gt_ptr loss_ptr
y_stride gt_stride loss_stride n_cols BLOCK_SIZE eps
in1 := y_ptr
in2 := gt_ptr
out := loss_ptr
B := BLOCK_SIZE
read1 := fun pid => pid * y_stride
read2 := fun pid => pid * gt_stride
write := fun pid => pid * loss_stride
mask := fun _pid j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
kldiv_opskldiv_backward_default_correctnessMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. Other reduction modes and host gradient scaling are not implied by these branch contracts. | Statement and preconditionsspecification kldiv_backward_default_correctness
(target_ptr new_grads_ptr : RegionName)
(target_stride new_grads_stride n_cols BLOCK_SIZE : Nat) :
kldivBackwardDefaultIO target_ptr new_grads_ptr
target_stride new_grads_stride n_cols BLOCK_SIZE
⊨ fun ts i => ts i * (0.0 - 1)
IO definition: kldivBackwardDefaultIOdef kldivBackwardDefaultIO (target_ptr new_grads_ptr : RegionName)
(target_stride new_grads_stride n_cols BLOCK_SIZE : Nat) : MaskedKernelIO₁ where
kernel := kldiv_backward_default target_ptr new_grads_ptr
target_stride new_grads_stride n_cols BLOCK_SIZE
inp := target_ptr
out := new_grads_ptr
B := BLOCK_SIZE
read := fun pid => pid * target_stride
write := fun pid => pid * new_grads_stride
mask := fun _pid j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
kldiv_opskldiv_backward_log_target_correctnessMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. Other reduction modes and host gradient scaling are not implied by these branch contracts. | Statement and preconditionsspecification kldiv_backward_log_target_correctness
(target_ptr new_grads_ptr : RegionName)
(target_stride new_grads_stride n_cols BLOCK_SIZE : Nat) :
kldivBackwardLogTargetIO target_ptr new_grads_ptr
target_stride new_grads_stride n_cols BLOCK_SIZE
⊨ fun ts i => 0.0 - Real.exp (ts i)
IO definition: kldivBackwardLogTargetIOdef kldivBackwardLogTargetIO (target_ptr new_grads_ptr : RegionName)
(target_stride new_grads_stride n_cols BLOCK_SIZE : Nat) : MaskedKernelIO₁ where
kernel := kldiv_backward_log_target target_ptr new_grads_ptr
target_stride new_grads_stride n_cols BLOCK_SIZE
inp := target_ptr
out := new_grads_ptr
B := BLOCK_SIZE
read := fun pid => pid * target_stride
write := fun pid => pid * new_grads_stride
mask := fun _pid j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
kldiv_tritonkldiv_backward_default_compute_correctMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. The IO headline covers only default non-reduced forward; other host reductions are external. | Statement and preconditionsspecification kldiv_backward_default_compute_correct
(input_ptr target_ptr : RegionName)
(input_stride target_stride n_cols BLOCK_SIZE : Nat)
(s : BlockState)
(hBS : 0 < BLOCK_SIZE)
(hLen : n_cols ≤ BLOCK_SIZE)
(hLenPos : 0 < n_cols)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => linearOffset s input_stride i)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := kldiv_backward_default input_ptr target_ptr
input_stride target_stride n_cols BLOCK_SIZE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < n_cols)
(fun i => (input_ptr, linearOffset s input_stride i)))
(expected := fun i => defaultSpec s target_ptr target_stride i)
Referenced model/specification definitions: Python source |
|---|
kldiv_tritonkldiv_backward_log_target_compute_correctMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. The IO headline covers only default non-reduced forward; other host reductions are external. | Statement and preconditionsspecification kldiv_backward_log_target_compute_correct
(input_ptr target_ptr : RegionName)
(input_stride target_stride n_cols BLOCK_SIZE : Nat)
(s : BlockState)
(hBS : 0 < BLOCK_SIZE)
(hLen : n_cols ≤ BLOCK_SIZE)
(hLenPos : 0 < n_cols)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => linearOffset s input_stride i)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := kldiv_backward_log_target input_ptr target_ptr
input_stride target_stride n_cols BLOCK_SIZE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < n_cols)
(fun i => (input_ptr, linearOffset s input_stride i)))
(expected := fun i => logTargetSpec s target_ptr target_stride i)
Referenced model/specification definitions: Python source |
|---|
kldiv_tritonkldiv_forward_default_none_compute_correctMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. The IO headline covers only default non-reduced forward; other host reductions are external. | Statement and preconditionsspecification kldiv_forward_default_none_compute_correct
(y_ptr gt_ptr loss_ptr : RegionName)
(y_stride gt_stride loss_stride n_cols BLOCK_SIZE : Nat)
(s : BlockState)
(hBS : 0 < BLOCK_SIZE)
(hLen : n_cols ≤ BLOCK_SIZE)
(hLenPos : 0 < n_cols)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => linearOffset s loss_stride i)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := kldiv_forward_default_none y_ptr gt_ptr loss_ptr
y_stride gt_stride loss_stride n_cols BLOCK_SIZE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < n_cols)
(fun i => (loss_ptr, linearOffset s loss_stride i)))
(expected := fun i =>
forwardDefaultSpec s y_ptr gt_ptr y_stride gt_stride i)
Referenced model/specification definitions: Python source |
|---|
kldiv_tritonkldiv_forward_log_target_none_compute_correctMathematical execution | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. The IO headline covers only default non-reduced forward; other host reductions are external. | Statement and preconditionsspecification kldiv_forward_log_target_none_compute_correct
(y_ptr gt_ptr loss_ptr : RegionName)
(y_stride gt_stride loss_stride n_cols BLOCK_SIZE : Nat)
(s : BlockState)
(hBS : 0 < BLOCK_SIZE)
(hLen : n_cols ≤ BLOCK_SIZE)
(hLenPos : 0 < n_cols)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => linearOffset s loss_stride i)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := kldiv_forward_log_target_none y_ptr gt_ptr loss_ptr
y_stride gt_stride loss_stride n_cols BLOCK_SIZE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < n_cols)
(fun i => (loss_ptr, linearOffset s loss_stride i)))
(expected := fun i =>
forwardLogTargetSpec s y_ptr gt_ptr y_stride gt_stride i)
Referenced model/specification definitions: Python source |
|---|
kldiv_tritonkldiv_forward_default_none_io_correctnessAbstract cast/store rounding | Configured model / stage Selected default/log-target backward branches and reduction=none forward branches. The IO headline covers only default non-reduced forward; other host reductions are external. | Statement and preconditionsspecification kldiv_forward_default_none_io_correctness (R : RoundingModel)
(y_ptr gt_ptr loss_ptr : RegionName)
(y_stride gt_stride loss_stride n_cols BLOCK_SIZE : Nat)
(hBS : 0 < BLOCK_SIZE)
(hly : loss_ptr ≠ y_ptr) (hlg : loss_ptr ≠ gt_ptr) :
kldivForwardDefaultNoneIO y_ptr gt_ptr loss_ptr y_stride gt_stride loss_stride
n_cols BLOCK_SIZE ⊨[R]
fun _ _ xs ys t j => klStreamSpec n_cols BLOCK_SIZE xs ys t j
IO definition: kldivForwardDefaultNoneIOdef kldivForwardDefaultNoneIO (y_ptr gt_ptr loss_ptr : RegionName)
(y_stride gt_stride loss_stride n_cols BLOCK_SIZE : Nat) :
StreamEmitMasked2DKernelIO₂ where
kernel := kldiv_forward_default_none y_ptr gt_ptr loss_ptr y_stride gt_stride
loss_stride n_cols BLOCK_SIZE
inp1 := y_ptr
inp2 := gt_ptr
out := loss_ptr
T := klNumSteps n_cols BLOCK_SIZE
B1 := BLOCK_SIZE
B2 := BLOCK_SIZE
C := BLOCK_SIZE
read1 := fun p₀ _ t j => p₀ * y_stride + (t.val * BLOCK_SIZE + j.val)
read2 := fun p₀ _ t j => p₀ * gt_stride + (t.val * BLOCK_SIZE + j.val)
write := fun p₀ _ t j => p₀ * loss_stride + (t.val * BLOCK_SIZE + j.val)
mask1 := fun _ _ t j => t.val * BLOCK_SIZE + j.val < n_cols
mask2 := fun _ _ t j => t.val * BLOCK_SIZE + j.val < n_cols
writeMask := fun _ _ t j => t.val * BLOCK_SIZE + j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
ksoftmax_tritonksoftmax_forward_plain_correctnessMathematical execution | Configured model / stage Plain forward softmax branch under positive depth and the declared strides. The Python backward and other flagged forward modes are outside this headline. | Statement and preconditionsspecification ksoftmax_forward_plain_correctness
(Y X : RegionName)
(stride_ym stride_yn stride_xm stride_xn K DEPTH : Nat)
(hD : 0 < DEPTH) :
ksoftmaxIO Y X stride_ym stride_yn stride_xm stride_xn K DEPTH ⊨
fun _ _ xs k => ksoftmaxSpec K DEPTH xs k
IO definition: ksoftmaxIOdef ksoftmaxIO (Y X : RegionName)
(stride_ym stride_yn stride_xm stride_xn K DEPTH : Nat) :
Masked2DKernelIO₁ where
kernel := ksoftmax_forward_plain Y X stride_ym stride_yn stride_xm
stride_xn K DEPTH
inp := X
out := Y
B := DEPTH
read := fun m n j => m * stride_xm + n * stride_xn + j.val
write := fun m n j => m * stride_ym + n * stride_yn + j.val
mask := fun _ _ j => j.val < K
Referenced model/specification definitions: Python source |
|---|
kv_cache_copykv_cache_copy_correctnessMathematical execution | Configured model / stage Sequence-length-one KV-cache x-block specialization with supplied block-table/context metadata; validity, cache bounds, and region separation are contract premises. | Statement and preconditionsspecification kv_cache_copy_correctness
(K V KCache VCache BLOCK_TABLES context_lengths : RegionName)
(SPLIT_X stride_kt stride_kh stride_kd
stride_vt stride_vh stride_vd
stride_kcb stride_kch stride_kcsplit_x stride_kcs
stride_vcb stride_vch stride_vcs stride_vcd
stride_bts stride_btb block_size HEAD_DIM KCACHE_X : Nat)
(hKV : KCache ≠ VCache) :
kvCacheCopyIO K V KCache VCache BLOCK_TABLES context_lengths
SPLIT_X stride_kt stride_kh stride_kd stride_vt stride_vh stride_vd
stride_kcb stride_kch stride_kcsplit_x stride_kcs stride_vcb stride_vch
stride_vcs stride_vcd stride_bts stride_btb block_size HEAD_DIM KCACHE_X
⊨ fun _ _ _ _ xs ys => (xs, ys)
IO definition: kvCacheCopyIOdef kvCacheCopyIO
(K V KCache VCache BLOCK_TABLES context_lengths : RegionName)
(SPLIT_X stride_kt stride_kh stride_kd
stride_vt stride_vh stride_vd
stride_kcb stride_kch stride_kcsplit_x stride_kcs
stride_vcb stride_vch stride_vcs stride_vcd
stride_bts stride_btb block_size HEAD_DIM KCACHE_X : Nat) : ChainMetaMasked2DKernelIO₂ₓ₂ where
kernel := copy_to_kvcache_seqlen1_xblock K V KCache VCache BLOCK_TABLES
context_lengths SPLIT_X stride_kt stride_kh stride_kd stride_vt stride_vh stride_vd
stride_kcb stride_kch stride_kcsplit_x stride_kcs stride_vcb stride_vch
stride_vcs stride_vcd stride_bts stride_btb block_size HEAD_DIM KCACHE_X
mbuf1 := context_lengths
mbuf2 := BLOCK_TABLES
in1 := K
in2 := V
out1 := KCache
out2 := VCache
B := KCACHE_X
mwin1 := fun pid₀ _ => pid₀
mwin2 := fun pid₀ _ m₁ =>
pid₀ * stride_bts + ((m₁ - 1) / block_size) * stride_btb
read1 := fun pid₀ pid₁ _ _ j =>
pid₀ * stride_kt + pid₁ * stride_kh +
(SPLIT_X * KCACHE_X + j.val) * stride_kd
read2 := fun pid₀ pid₁ _ _ j =>
pid₀ * stride_vt + pid₁ * stride_vh +
(SPLIT_X * KCACHE_X + j.val) * stride_vd
mask1 := fun _ _ _ _ j => SPLIT_X * KCACHE_X + j.val < HEAD_DIM
mask2 := fun _ _ _ _ j => SPLIT_X * KCACHE_X + j.val < HEAD_DIM
write1 := fun _ pid₁ m₁ m₂ j =>
m₂ * stride_kcb + pid₁ * stride_kch + SPLIT_X * stride_kcsplit_x +
((m₁ - 1) % block_size) * stride_kcs + j.val
write2 := fun _ pid₁ m₁ m₂ j =>
m₂ * stride_vcb + pid₁ * stride_vch +
((m₁ - 1) % block_size) * stride_vcs +
(SPLIT_X * KCACHE_X + j.val) * stride_vcd
Referenced model/specification definitions: Python source |
|---|
kv_cache_fillingfill_kv_cache_output_summary_generalMathematical execution | Configured model / stage Unquantized K/V tile-copy facts under selected cache metadata and layouts; quantized paths must not be inferred from projection/copy facts. | Statement and preconditionsspecification fill_kv_cache_output_summary_general
(KStates VStates KCaches VCaches QStartLoc QSeqLens KVSeqLens
BlockOffsets : RegionName)
(SIDX BIDX KV_BLOCK_IDX
num_heads head_dim head_dim_v stride_kss stride_ksh stride_ksd
stride_vss stride_vsh stride_vsd stride_kcn stride_kcb stride_kch
stride_kcd stride_vcn stride_vcb stride_vch stride_vcd stride_boff
BLOCK BLOCK_D BLOCK_DV BLOCK_H : Nat)
(s : BlockState)
(hKInj : Function.Injective
(fun idx : TileIndex [BLOCK_H, BLOCK_D] =>
kCacheOffset s BlockOffsets BIDX KV_BLOCK_IDX stride_kcn stride_kcb
stride_kch stride_kcd stride_boff idx))
(hVInj : Function.Injective
(fun idx : TileIndex [BLOCK_H, BLOCK_DV] =>
vCacheOffset s BlockOffsets BIDX KV_BLOCK_IDX stride_vcn stride_vcb
stride_vch stride_vcd stride_boff idx)) :
(∃ alg, (fill_kv_cache_kernel_surface KStates VStates KCaches VCaches
QStartLoc QSeqLens KVSeqLens BlockOffsets num_heads head_dim head_dim_v
stride_kss stride_ksh stride_ksd stride_vss stride_vsh stride_vsd
stride_kcn stride_kcb stride_kch stride_kcd stride_vcn stride_vcb
stride_vch stride_vcd stride_boff BLOCK BLOCK_D BLOCK_DV
BLOCK_H).toAlgorithm? = Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := fill_k_cache_tile KStates KCaches BlockOffsets
SIDX BIDX KV_BLOCK_IDX stride_kss stride_ksh stride_ksd
stride_kcn stride_kcb stride_kch stride_kcd stride_boff
num_heads head_dim BLOCK_H BLOCK_D)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s num_heads head_dim BLOCK_H BLOCK_D)
(fun idx => (KCaches,
kCacheOffset s BlockOffsets BIDX KV_BLOCK_IDX stride_kcn stride_kcb
stride_kch stride_kcd stride_boff idx)))
(expected := fun idx =>
s.readMem KStates
(kSourceOffset s SIDX stride_kss stride_ksh stride_ksd idx))) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := fill_v_cache_tile VStates VCaches BlockOffsets
SIDX BIDX KV_BLOCK_IDX stride_vss stride_vsh stride_vsd
stride_vcn stride_vcb stride_vch stride_vcd stride_boff
num_heads head_dim_v BLOCK_H BLOCK_DV)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s num_heads head_dim_v BLOCK_H BLOCK_DV)
(fun idx => (VCaches,
vCacheOffset s BlockOffsets BIDX KV_BLOCK_IDX stride_vcn stride_vcb
stride_vch stride_vcd stride_boff idx)))
(expected := fun idx =>
s.readMem VStates
(vSourceOffset s SIDX stride_vss stride_vsh stride_vsd idx)))
Referenced model/specification definitions: Python source |
|---|
kv_cache_fillingfill_quant_meta_store_io_correctnessMathematical execution | Precomputed-input slice Stores supplied quantization metadata MetaPre; it does not compute int8/int4 scales, zero points, or quantized cache values. | Statement and preconditionsspecification fill_quant_meta_store_io_correctness
(MetaPre MetaOut BlockOffsets : RegionName)
(BIDX KV_BLOCK_IDX SZD stride_mn stride_mb stride_mh stride_md stride_boff
num_heads BLOCK_H : Nat)
(hOutInj : ∀ m : Nat, Function.Injective
(fun i : Fin BLOCK_H =>
m * stride_mn + BIDX * stride_mb + i.val * stride_mh
+ SZD * stride_md)) :
metaStoreIO MetaPre MetaOut BlockOffsets BIDX KV_BLOCK_IDX SZD stride_mn
stride_mb stride_mh stride_md stride_boff num_heads BLOCK_H
⊨ fun _pid _m xs i => xs i
IO definition: metaStoreIOnoncomputable def metaStoreIO (MetaPre MetaOut BlockOffsets : RegionName)
(BIDX KV_BLOCK_IDX SZD stride_mn stride_mb stride_mh stride_md stride_boff
num_heads BLOCK_H : Nat) : Meta1MaskedTileKernelIO₁ where
kernel := fill_quant_meta_store_slice MetaPre MetaOut BlockOffsets BIDX
KV_BLOCK_IDX SZD stride_mn stride_mb stride_mh stride_md stride_boff
num_heads BLOCK_H
mbuf := BlockOffsets
inp := MetaPre
out := MetaOut
shape := [BLOCK_H]
mwin := fun pid => pid * stride_boff + KV_BLOCK_IDX
read := fun _pid _m i => i.1.val
write := fun _pid m i =>
m * stride_mn + BIDX * stride_mb + i.1.val * stride_mh + SZD * stride_md
mask := fun _pid _m i => i.1.val < num_heads
Referenced model/specification definitions: Python source |
|---|
l2_norm_bwdl2_norm_bwd_kernel_correctnessMathematical execution | Configured model / stage One L2-normalization backward row under the declared mask/stride/epsilon model. Mathematical sqrt and division do not certify concrete hardware exceptional values. | Statement and preconditionsspecification l2_norm_bwd_kernel_correctness
(X DY DX : RegionName)
(stride_x_row N : Nat) (eps : ℝ) (BLOCK_N : Nat) :
l2BwdIO X DY DX stride_x_row N eps BLOCK_N
⊨ fun xs dys i => l2BwdSpec N BLOCK_N eps xs dys i
IO definition: l2BwdIOdef l2BwdIO (X DY DX : RegionName)
(stride_x_row N : Nat) (eps : ℝ) (BLOCK_N : Nat) : MaskedKernelIO₂ where
kernel := l2_norm_bwd_kernel X DY DX stride_x_row N eps BLOCK_N
in1 := X
in2 := DY
out := DX
B := BLOCK_N
read1 := fun pid => pid * stride_x_row
read2 := fun pid => pid * stride_x_row
write := fun pid => pid * stride_x_row
mask := fun _ j => j.val < N
Referenced model/specification definitions: Python source |
|---|
l2_norm_triton1l2_norm_fwd_1pass_kernel_correctnessMathematical execution | Configured model / stage One-pass L2-normalization forward row with explicit row stride, mask, and epsilon. This is a mathematical kernel contract, not a host/tensor-layout certification. | Statement and preconditionsspecification l2_norm_fwd_1pass_kernel_correctness
(X Y : RegionName)
(stride_x_row N : Nat) (eps : ℝ) (BLOCK_N : Nat) :
l2NormIO X Y stride_x_row N eps BLOCK_N ⊨
fun xs i => l2NormSpec N BLOCK_N eps xs i
IO definition: l2NormIOdef l2NormIO (X Y : RegionName)
(stride_x_row N : Nat) (eps : ℝ) (BLOCK_N : Nat) :
MaskedKernelIO₁ where
kernel := l2_norm_fwd_1pass_kernel X Y stride_x_row N eps BLOCK_N
inp := X
out := Y
B := BLOCK_N
read := fun pid => pid * stride_x_row
write := fun pid => pid * stride_x_row
mask := fun _ j => j.val < N
Referenced model/specification definitions: Python source |
|---|
l2_norm_triton2l2_norm_fwd_1pass_kernel_correctnessMathematical execution | Configured model / stage Separate one-pass forward and backward L2-normalization rows with their own IO targets and mathematical sqrt/division semantics. | Statement and preconditionsspecification l2_norm_fwd_1pass_kernel_correctness
(X Y : RegionName)
(stride_x_row N : Nat) (eps : ℝ) (BLOCK_N : Nat) :
l2FwdIO X Y stride_x_row N eps BLOCK_N ⊨
fun xs i => l2FwdSpec N BLOCK_N eps xs i
IO definition: l2FwdIOdef l2FwdIO (X Y : RegionName)
(stride_x_row N : Nat) (eps : ℝ) (BLOCK_N : Nat) :
MaskedKernelIO₁ where
kernel := l2_norm_fwd_1pass_kernel X Y stride_x_row N eps BLOCK_N
inp := X
out := Y
B := BLOCK_N
read := fun pid => pid * stride_x_row
write := fun pid => pid * stride_x_row
mask := fun _ j => j.val < N
Referenced model/specification definitions: Python source |
|---|
l2_norm_triton2l2_norm_bwd_kernel_correctnessMathematical execution | Configured model / stage Separate one-pass forward and backward L2-normalization rows with their own IO targets and mathematical sqrt/division semantics. | Statement and preconditionsspecification l2_norm_bwd_kernel_correctness
(X DY DX : RegionName)
(stride_x_row N : Nat) (eps : ℝ) (BLOCK_N : Nat) :
l2BwdIO X DY DX stride_x_row N eps BLOCK_N ⊨
fun xs dys i => l2BwdSpec N BLOCK_N eps xs dys i
IO definition: l2BwdIOdef l2BwdIO (X DY DX : RegionName)
(stride_x_row N : Nat) (eps : ℝ) (BLOCK_N : Nat) :
MaskedKernelIO₂ where
kernel := l2_norm_bwd_kernel X DY DX stride_x_row N eps BLOCK_N
in1 := X
in2 := DY
out := DX
B := BLOCK_N
read1 := fun pid => pid * stride_x_row
read2 := fun pid => pid * stride_x_row
write := fun pid => pid * stride_x_row
mask := fun _ j => j.val < N
Referenced model/specification definitions: Python source |
|---|
layer_norm_ligerlayer_norm_liger_forward_output_summaryMathematical execution | Configured model / stage Forward value facts for Y/mean/RSTD; the Python backward kernel is outside these headlines. | Statement and preconditionsspecification layer_norm_liger_forward_output_summary
(Y X W B Mean RSTD : RegionName)
(Y_row_stride X_row_stride W_row_stride B_row_stride
Mean_row_stride RSTD_row_stride n_cols BLOCK_SIZE : Nat)
(eps : ℝ) (s : BlockState)
(hYMean : Y ≠ Mean) (hYRSTD : Y ≠ RSTD)
(hMeanY : Mean ≠ Y) (hMeanRSTD : Mean ≠ RSTD)
(hRSTDY : RSTD ≠ Y)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => yOffset s Y_row_stride i)) :
(∃ alg, (layer_norm_liger_forward_surface Y Y_row_stride X
X_row_stride W W_row_stride B B_row_stride Mean Mean_row_stride RSTD
RSTD_row_stride n_cols eps BLOCK_SIZE).toAlgorithm? = Except.ok alg) ∧
((ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_liger_forward Y X W B Mean RSTD Y_row_stride
X_row_stride Mean_row_stride RSTD_row_stride n_cols BLOCK_SIZE eps)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < n_cols)
(fun i => (Y, yOffset s Y_row_stride i)))
(expected := fun i =>
layernormYSpec s X W B X_row_stride n_cols BLOCK_SIZE eps i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_liger_forward Y X W B Mean RSTD Y_row_stride
X_row_stride Mean_row_stride RSTD_row_stride n_cols BLOCK_SIZE eps)
(initialState := s)
(write := fun _ : PUnit => some (Mean, s.pid * Mean_row_stride))
(expected := fun _ =>
WithBot.unbotD 0
(layernormMeanCarrier s X X_row_stride n_cols BLOCK_SIZE))) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_liger_forward Y X W B Mean RSTD Y_row_stride
X_row_stride Mean_row_stride RSTD_row_stride n_cols BLOCK_SIZE eps)
(initialState := s)
(write := fun _ : PUnit => some (RSTD, s.pid * RSTD_row_stride))
(expected := fun _ =>
WithBot.unbotD 0
(layernormInvVarCarrier s X X_row_stride n_cols BLOCK_SIZE eps))))
Referenced model/specification definitions: Python source |
|---|
layer_norm_ligerlayer_norm_liger_scalar_stores_io_correctnessMathematical execution | Precomputed-input slice Stores precomputed MeanPre/RSTDPre scalars; these IO kernels do not reduce or normalize the input row. | Statement and preconditionsspecification layer_norm_liger_scalar_stores_io_correctness
(MeanPre Mean RSTDPre RSTD : RegionName)
(Mean_row_stride RSTD_row_stride : Nat) :
(mean_storeIO MeanPre Mean Mean_row_stride
⊨ fun _pid xs _ => xs PUnit.unit) ∧
(rstd_storeIO RSTDPre RSTD RSTD_row_stride
⊨ fun _pid xs _ => xs PUnit.unit)
IO definition: mean_storeIOdef mean_storeIO (MeanPre Mean : RegionName) (Mean_row_stride : Nat) : MaskedTileKernelIO₁ where
kernel := layer_norm_liger_forward_mean_store_slice MeanPre Mean Mean_row_stride
inp := MeanPre
out := Mean
shape := []
read := fun pid _ => pid * Mean_row_stride
write := fun pid _ => pid * Mean_row_stride
mask := fun _pid _ => True
IO definition: rstd_storeIOdef rstd_storeIO (RSTDPre RSTD : RegionName) (RSTD_row_stride : Nat) : MaskedTileKernelIO₁ where
kernel := layer_norm_liger_forward_rstd_store_slice RSTDPre RSTD RSTD_row_stride
inp := RSTDPre
out := RSTD
shape := []
read := fun pid _ => pid * RSTD_row_stride
write := fun pid _ => pid * RSTD_row_stride
mask := fun _pid _ => True
Referenced model/specification definitions: Python source |
|---|
layer_norm_ligerlayer_norm_liger_scalar_stores_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding contracts for the precomputed mean/RSTD scalar stores. | Statement and preconditionsspecification layer_norm_liger_scalar_stores_io_correctnessR (R : RoundingModel)
(MeanPre Mean RSTDPre RSTD : RegionName)
(Mean_row_stride RSTD_row_stride : Nat) :
(mean_storeIO MeanPre Mean Mean_row_stride
⊨[R, FloatDType.real] fun _pid xs _ => xs PUnit.unit) ∧
(rstd_storeIO RSTDPre RSTD RSTD_row_stride
⊨[R, FloatDType.real] fun _pid xs _ => xs PUnit.unit)
IO definition: mean_storeIOdef mean_storeIO (MeanPre Mean : RegionName) (Mean_row_stride : Nat) : MaskedTileKernelIO₁ where
kernel := layer_norm_liger_forward_mean_store_slice MeanPre Mean Mean_row_stride
inp := MeanPre
out := Mean
shape := []
read := fun pid _ => pid * Mean_row_stride
write := fun pid _ => pid * Mean_row_stride
mask := fun _pid _ => True
IO definition: rstd_storeIOdef rstd_storeIO (RSTDPre RSTD : RegionName) (RSTD_row_stride : Nat) : MaskedTileKernelIO₁ where
kernel := layer_norm_liger_forward_rstd_store_slice RSTDPre RSTD RSTD_row_stride
inp := RSTDPre
out := RSTD
shape := []
read := fun pid _ => pid * RSTD_row_stride
write := fun pid _ => pid * RSTD_row_stride
mask := fun _pid _ => True
Referenced model/specification definitions: Python source |
|---|
layer_norm_opslayer_norm_ops_fwd_plain_bias_all_outputs_compute_correct_generalMathematical execution | Precomputed-input slice Plain forward headline connects stores of supplied YPre/MeanPre/RstdPre; their producers are outside these store kernels. | Statement and preconditionsspecification layer_norm_ops_fwd_plain_bias_all_outputs_compute_correct_general
(ValuePre MeanPre RstdPre Y Mean Rstd : RegionName) (s : BlockState)
(stride_y_row N BLOCK_N : Nat) :
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_fwd_y_store_slice ValuePre Y
stride_y_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (Y, fwdYOffset s stride_y_row i)))
(expected := fun i : Fin BLOCK_N =>
fwdYStoreSpec s ValuePre stride_y_row i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_fwd_mean_store_slice MeanPre Mean)
(initialState := s)
(write := fun _ : PUnit => some (Mean, meanRowOffset s))
(expected := fun _ => meanStoreSpec s MeanPre)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_fwd_rstd_store_slice RstdPre Rstd)
(initialState := s)
(write := fun _ : PUnit => some (Rstd, meanRowOffset s))
(expected := fun _ => rstdStoreSpec s RstdPre))
Referenced model/specification definitions: Python source |
|---|
layer_norm_opslayer_norm_ops_fwd_rms_bias_all_outputs_compute_correct_generalMathematical execution | Precomputed-input slice RMS forward headline connects supplied YPre/RstdPre stores, not the complete normalization producer. | Statement and preconditionsspecification layer_norm_ops_fwd_rms_bias_all_outputs_compute_correct_general
(ValuePre RstdPre Y Rstd : RegionName) (s : BlockState)
(stride_y_row N BLOCK_N : Nat) :
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_fwd_y_store_slice ValuePre Y
stride_y_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (Y, fwdYOffset s stride_y_row i)))
(expected := fun i : Fin BLOCK_N =>
fwdYStoreSpec s ValuePre stride_y_row i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_fwd_rstd_store_slice RstdPre Rstd)
(initialState := s)
(write := fun _ : PUnit => some (Rstd, meanRowOffset s))
(expected := fun _ => rstdStoreSpec s RstdPre))
Referenced model/specification definitions: Python source |
|---|
layer_norm_opslayer_norm_ops_fwd_residual_bias_all_outputs_compute_correct_generalMathematical execution | Precomputed-input slice Residual forward headline concerns supplied output/statistic store slices; it does not compose the full residual/normalization producer. | Statement and preconditionsspecification layer_norm_ops_fwd_residual_bias_all_outputs_compute_correct_general
(ResidualPre ValuePre MeanPre RstdPre RESIDUAL_OUT Y Mean Rstd : RegionName)
(s : BlockState)
(stride_res_out_row stride_y_row N BLOCK_N : Nat) :
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_fwd_residual_out_store_slice ResidualPre
RESIDUAL_OUT stride_res_out_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (RESIDUAL_OUT, fwdResidualOutOffset s stride_res_out_row i)))
(expected := fun i : Fin BLOCK_N =>
fwdResidualOutStoreSpec s ResidualPre stride_res_out_row i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_fwd_y_store_slice ValuePre Y
stride_y_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (Y, fwdYOffset s stride_y_row i)))
(expected := fun i : Fin BLOCK_N =>
fwdYStoreSpec s ValuePre stride_y_row i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_fwd_mean_store_slice MeanPre Mean)
(initialState := s)
(write := fun _ : PUnit => some (Mean, meanRowOffset s))
(expected := fun _ => meanStoreSpec s MeanPre)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_fwd_rstd_store_slice RstdPre Rstd)
(initialState := s)
(write := fun _ : PUnit => some (Rstd, meanRowOffset s))
(expected := fun _ => rstdStoreSpec s RstdPre))
Referenced model/specification definitions: Python source |
|---|
layer_norm_opslayer_norm_ops_bwd_rms_core_outputs_compute_correct_generalMathematical execution | Configured model / stage Backward RMS one-row DX/DW and reduction-stage facts with input Rstd; not a whole multi-row parameter-gradient reduction. | Statement and preconditionsspecification layer_norm_ops_bwd_rms_core_outputs_compute_correct_general
(X Xhat W DY Rstd C1 DX DW : RegionName) (s : BlockState)
(stride_xhat_row stride_dy_row stride_x_row stride_dx_row N BLOCK_N : Nat)
(hDWDX : DW ≠ DX) :
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_bwd_c1_reduction_slice Xhat W DY C1
stride_xhat_row stride_dy_row N BLOCK_N)
(initialState := s)
(write := fun _ : PUnit => some (C1, s.pid))
(expected := fun _ : PUnit =>
bwdC1ReductionSpec s Xhat W DY stride_xhat_row stride_dy_row N BLOCK_N)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_bwd_rms_dx_from_c1_slice Xhat W DY Rstd C1 DX
stride_xhat_row stride_dy_row stride_dx_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (DX, bwdRmsDXOffset s stride_dx_row i)))
(expected := fun i : Fin BLOCK_N =>
bwdRmsDXFromC1Spec s Xhat W DY Rstd C1 stride_xhat_row stride_dy_row i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_bwd_rms_one_row X W DY DX DW Rstd
stride_x_row stride_dy_row stride_dx_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (DW, bwdRmsDWOffset s N i)))
(expected := fun i : Fin BLOCK_N =>
bwdRmsDWSpec s X DY Rstd stride_x_row stride_dy_row N BLOCK_N i))
Referenced model/specification definitions: Python source |
|---|
layer_norm_opslayer_norm_ops_bwd_plain_bias_core_outputs_compute_correct_generalMathematical execution | Configured model / stage Backward plain-bias one-row DX/DW/DB and C1/C2 reduction facts; host aggregation of parameter gradients is outside the statement. | Statement and preconditionsspecification layer_norm_ops_bwd_plain_bias_core_outputs_compute_correct_general
(X Xhat W DY DX DW DB Mean Rstd C1 C2 : RegionName) (s : BlockState)
(stride_xhat_row stride_dy_row stride_x_row stride_dx_row N BLOCK_N : Nat)
(hDWDX : DW ≠ DX) (hDWDB : DW ≠ DB)
(hDBDX : DB ≠ DX) (hDBDW : DB ≠ DW) :
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_bwd_c1_reduction_slice Xhat W DY C1
stride_xhat_row stride_dy_row N BLOCK_N)
(initialState := s)
(write := fun _ : PUnit => some (C1, s.pid))
(expected := fun _ : PUnit =>
bwdC1ReductionSpec s Xhat W DY stride_xhat_row stride_dy_row N BLOCK_N)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_bwd_c2_reduction_slice W DY C2
stride_dy_row N BLOCK_N)
(initialState := s)
(write := fun _ : PUnit => some (C2, s.pid))
(expected := fun _ : PUnit =>
bwdC2ReductionSpec s W DY stride_dy_row N BLOCK_N)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_bwd_plain_dx_from_c1_c2_slice Xhat W DY Rstd
C1 C2 DX stride_xhat_row stride_dy_row stride_dx_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (DX, bwdRmsDXOffset s stride_dx_row i)))
(expected := fun i : Fin BLOCK_N =>
bwdPlainDXFromC1C2Spec s Xhat W DY Rstd C1 C2 stride_xhat_row
stride_dy_row i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_bwd_plain_bias_one_row X W DY DX DW DB Mean Rstd
stride_x_row stride_dy_row stride_dx_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (DW, bwdParamGradOffset s N i)))
(expected := fun i : Fin BLOCK_N =>
bwdPlainBiasDWSpec s X DY Mean Rstd stride_x_row stride_dy_row N
BLOCK_N i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_bwd_plain_bias_one_row X W DY DX DW DB Mean Rstd
stride_x_row stride_dy_row stride_dx_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (DB, bwdParamGradOffset s N i)))
(expected := fun i : Fin BLOCK_N =>
bwdBiasDBSpec s DY stride_dy_row i))
Referenced model/specification definitions: Python source |
|---|
layer_norm_opslayer_norm_ops_bwd_residual_add_all_outputs_compute_correct_generalMathematical execution | Precomputed-input slice Adds supplied DXBase and residual-gradient values in a store slice; upstream normalization-gradient production is external. | Statement and preconditionsspecification layer_norm_ops_bwd_residual_add_all_outputs_compute_correct_general
(DXBase DRESIDUAL DX DRESIDUAL_IN : RegionName) (s : BlockState)
(stride_dx_row stride_dres_row stride_dres_in_row N BLOCK_N : Nat)
(hDXDresIn : DX ≠ DRESIDUAL_IN)
(hDresInDX : DRESIDUAL_IN ≠ DX) :
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_bwd_residual_add_store_slice DXBase DRESIDUAL
DX DRESIDUAL_IN stride_dx_row stride_dres_row stride_dres_in_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (DX, bwdRmsDXOffset s stride_dx_row i)))
(expected := fun i : Fin BLOCK_N =>
bwdResidualAddSpec s DXBase DRESIDUAL stride_dx_row stride_dres_row i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_ops_bwd_residual_add_store_slice DXBase DRESIDUAL
DX DRESIDUAL_IN stride_dx_row stride_dres_row stride_dres_in_row N BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => i.val < N)
(fun i => (DRESIDUAL_IN, bwdDResidualInOffset s stride_dres_in_row i)))
(expected := fun i : Fin BLOCK_N =>
bwdResidualAddSpec s DXBase DRESIDUAL stride_dx_row stride_dres_row i))
Referenced model/specification definitions: Python source |
|---|
layer_norm_opslayer_norm_ops_scalar_stores_io_correctnessMathematical execution | Precomputed-input slice Copies supplied MeanPre/RstdPre scalars to the output statistic buffers. | Statement and preconditionsspecification layer_norm_ops_scalar_stores_io_correctness
(MeanPre Mean RstdPre Rstd : RegionName) :
(fwd_mean_storeIO MeanPre Mean ⊨ fun _pid xs _ => xs PUnit.unit) ∧
(fwd_rstd_storeIO RstdPre Rstd ⊨ fun _pid xs _ => xs PUnit.unit)
IO definition: fwd_mean_storeIOdef fwd_mean_storeIO (MeanPre Mean : RegionName) : MaskedTileKernelIO₁ where
kernel := layer_norm_ops_fwd_mean_store_slice MeanPre Mean
inp := MeanPre
out := Mean
shape := []
read := fun pid _ => pid
write := fun pid _ => pid
mask := fun _pid _ => True
IO definition: fwd_rstd_storeIOdef fwd_rstd_storeIO (RstdPre Rstd : RegionName) : MaskedTileKernelIO₁ where
kernel := layer_norm_ops_fwd_rstd_store_slice RstdPre Rstd
inp := RstdPre
out := Rstd
shape := []
read := fun pid _ => pid
write := fun pid _ => pid
mask := fun _pid _ => True
Referenced model/specification definitions: Python source |
|---|
layer_norm_opslayer_norm_ops_scalar_stores_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding version of the supplied mean/Rstd scalar-store contracts. | Statement and preconditionsspecification layer_norm_ops_scalar_stores_io_correctnessR (R : RoundingModel)
(MeanPre Mean RstdPre Rstd : RegionName) :
(fwd_mean_storeIO MeanPre Mean
⊨[R, FloatDType.real] fun _pid xs _ => xs PUnit.unit) ∧
(fwd_rstd_storeIO RstdPre Rstd
⊨[R, FloatDType.real] fun _pid xs _ => xs PUnit.unit)
IO definition: fwd_mean_storeIOdef fwd_mean_storeIO (MeanPre Mean : RegionName) : MaskedTileKernelIO₁ where
kernel := layer_norm_ops_fwd_mean_store_slice MeanPre Mean
inp := MeanPre
out := Mean
shape := []
read := fun pid _ => pid
write := fun pid _ => pid
mask := fun _pid _ => True
IO definition: fwd_rstd_storeIOdef fwd_rstd_storeIO (RstdPre Rstd : RegionName) : MaskedTileKernelIO₁ where
kernel := layer_norm_ops_fwd_rstd_store_slice RstdPre Rstd
inp := RstdPre
out := Rstd
shape := []
read := fun pid _ => pid
write := fun pid _ => pid
mask := fun _pid _ => True
Referenced model/specification definitions: Python source |
|---|
layer_norm_welfoldlayer_norm_welfold_output_summary_generalMathematical execution | Configured model / stage Separate reduction, normalization, and linked-surface facts with explicit intermediate-state premises. | Statement and preconditionsspecification layer_norm_welfold_output_summary_general
(in_out_ptr0 in_out_ptr1 in_ptr0 in_ptr1 in_ptr2 out_ptr0 : RegionName)
(xnumel rnumel XBLOCK RBLOCK : Nat) (hLe : rnumel ≤ RBLOCK)
(s : BlockState)
(hMeanRstd : in_out_ptr0 ≠ in_out_ptr1)
(hInMean : in_ptr0 ≠ in_out_ptr0) :
-- (1) the full faithful surface lowers to the algorithm layer
(∃ alg,
(layer_norm_welfold_surface in_out_ptr0 in_out_ptr1 in_ptr0 in_ptr1
in_ptr2 out_ptr0 xnumel rnumel XBLOCK RBLOCK).toAlgorithm?
= Except.ok alg) ∧
-- (2) reduction phases: genuine mean and rstd, end-to-end from `in_ptr0`
((ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_welfold_reduce_slice in_out_ptr0 in_out_ptr1
in_ptr0 rnumel RBLOCK)
(initialState := s)
(write := fun _ : PUnit => some (in_out_ptr0, s.pids 0))
(expected := fun _ => rowMeanSpec s in_ptr0 rnumel)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_welfold_reduce_slice in_out_ptr0 in_out_ptr1
in_ptr0 rnumel RBLOCK)
(initialState := s)
(write := fun _ : PUnit => some (in_out_ptr1, s.pids 0))
(expected := fun _ => rowRstdSpec s in_ptr0 rnumel))) ∧
-- (3) normalize face: genuine `Y` chunk for every mean/rstd-genuine state
(∀ s' : BlockState,
s'.readMem in_out_ptr0 (s'.pids 0) = rowMeanSpec s' in_ptr0 rnumel →
s'.readMem in_out_ptr1 (s'.pids 0) = rowRstdSpec s' in_ptr0 rnumel →
ComputeCorrect.Realizes_without_Rounding
(kernel := layer_norm_welfold_normalize_slice in_out_ptr0 in_out_ptr1
in_ptr0 in_ptr1 in_ptr2 out_ptr0 rnumel RBLOCK)
(initialState := s')
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin RBLOCK => activeLane s' rnumel RBLOCK i)
(fun i => (out_ptr0, yOffset s' rnumel RBLOCK i)))
(expected := fun i : Fin RBLOCK =>
rowYSpec s' in_ptr0 in_ptr1 in_ptr2 rnumel (rIndex s' RBLOCK i)))
Referenced model/specification definitions: Python source |
|---|
layer_norm_welfoldlayer_norm_welfold_normalize_io_correctnessMathematical execution | Precomputed-input slice Normalizes using supplied mean and inverse variance; this IO kernel does not compute those statistics. | Statement and preconditionsspecification layer_norm_welfold_normalize_io_correctness
(in_out_ptr0 in_out_ptr1 in_ptr0 in_ptr1 in_ptr2 out_ptr0 : RegionName)
(rnumel RBLOCK : Nat) :
normalizeIO in_out_ptr0 in_out_ptr1 in_ptr0 in_ptr1 in_ptr2 out_ptr0 rnumel
RBLOCK
⊨ fun _p₀ _p₁ m1 m2 xs ws bs i => ((xs i - m1) * m2) * ws i + bs i
IO definition: normalizeIOdef normalizeIO
(in_out_ptr0 in_out_ptr1 in_ptr0 in_ptr1 in_ptr2 out_ptr0 : RegionName)
(rnumel RBLOCK : Nat) : Scalar2Tile3KernelIO where
kernel := layer_norm_welfold_normalize_slice in_out_ptr0 in_out_ptr1 in_ptr0
in_ptr1 in_ptr2 out_ptr0 rnumel RBLOCK
sbuf1 := in_out_ptr0
sbuf2 := in_out_ptr1
tbuf1 := in_ptr0
tbuf2 := in_ptr1
tbuf3 := in_ptr2
out := out_ptr0
shape := [RBLOCK]
swin1 := fun p₀ _p₁ => p₀
swin2 := fun p₀ _p₁ => p₀
read1 := fun p₀ p₁ i => p₁ * RBLOCK + i.1.val + rnumel * p₀
read2 := fun _p₀ p₁ i => p₁ * RBLOCK + i.1.val
read3 := fun _p₀ p₁ i => p₁ * RBLOCK + i.1.val
write := fun p₀ p₁ i => p₁ * RBLOCK + i.1.val + rnumel * p₀
mask := fun _p₀ p₁ i => p₁ * RBLOCK + i.1.val < rnumel
Referenced model/specification definitions: Python source |
|---|
layernorm_fwd_tritonlayernorm_fwd_triton_output_summaryMathematical execution | Configured model / stage Forward strided full-row normalization for the modeled stream extent with positive block size and input/output separation. No backward or host-layout composition claim. | Statement and preconditionsspecification layernorm_fwd_triton_output_summary
(X W Y : RegionName)
(stride_x_N stride_x_hn stride_x_hd
stride_y_N stride_y_hn stride_y_hd stride_w_hn stride_w_hd
N BLOCK_SIZE : Nat)
(eps : ℝ) (s : BlockState)
(hBlockPos : 0 < BLOCK_SIZE)
(hXYNe : X ≠ Y)
(hWYNe : W ≠ Y) :
(∃ alg, (layernorm_fwd_triton X W Y
stride_x_N stride_x_hn stride_x_hd
stride_y_N stride_y_hn stride_y_hd stride_w_hn stride_w_hd
N BLOCK_SIZE eps).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := layernorm_fwd_triton X W Y
stride_x_N stride_x_hn stride_x_hd
stride_y_N stride_y_hn stride_y_hd stride_w_hn stride_w_hd
N BLOCK_SIZE eps)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun _ : Fin N => True)
(fun i => (Y, yColOffset s stride_y_N stride_y_hn i.val)))
(expected := fun i =>
layernormYFullNSpec s X W stride_x_N stride_x_hn stride_w_hn
N BLOCK_SIZE eps i)
Referenced model/specification definitions: Python source |
|---|
layernorm_fwd_tritonlayernorm_fwd_triton_io_correctnessAbstract cast/store rounding | Configured model / stage Forward strided full-row normalization for the modeled stream extent with positive block size and input/output separation. No backward or host-layout composition claim. | Statement and preconditionsspecification layernorm_fwd_triton_io_correctness (R : RoundingModel)
(X W Y : RegionName)
(stride_x_N stride_x_hn stride_x_hd
stride_y_N stride_y_hn stride_y_hd stride_w_hn stride_w_hd
N BLOCK_SIZE : Nat) (eps : ℝ)
(hBlockPos : 0 < BLOCK_SIZE) (hXYNe : X ≠ Y) (hWYNe : W ≠ Y) :
layernormKernelIO X W Y stride_x_N stride_x_hn stride_x_hd
stride_y_N stride_y_hn stride_y_hd stride_w_hn stride_w_hd
N BLOCK_SIZE eps ⊨[R]
fun _ _ xs ws t j => lnStreamSpec N BLOCK_SIZE eps xs ws t j
IO definition: layernormKernelIOdef layernormKernelIO (X W Y : RegionName)
(stride_x_N stride_x_hn stride_x_hd
stride_y_N stride_y_hn stride_y_hd
stride_w_hn stride_w_hd N BLOCK_SIZE : Nat) (eps : ℝ) :
StreamEmitMasked2DKernelIO₂ where
kernel := layernorm_fwd_triton X W Y stride_x_N stride_x_hn stride_x_hd
stride_y_N stride_y_hn stride_y_hd stride_w_hn stride_w_hd N BLOCK_SIZE eps
inp1 := X
inp2 := W
out := Y
T := lnNumSteps N BLOCK_SIZE
B1 := BLOCK_SIZE
B2 := BLOCK_SIZE
C := BLOCK_SIZE
read1 := fun p₀ p₁ t j =>
p₀ * stride_x_N + p₁ * stride_x_hn + (t.val * BLOCK_SIZE + j.val)
read2 := fun _ p₁ t j => p₁ * stride_w_hn + (t.val * BLOCK_SIZE + j.val)
write := fun p₀ p₁ t j =>
p₀ * stride_y_N + p₁ * stride_y_hn + (t.val * BLOCK_SIZE + j.val)
mask1 := fun _ _ t j => t.val * BLOCK_SIZE + j.val < N
mask2 := fun _ _ t j => t.val * BLOCK_SIZE + j.val < N
writeMask := fun _ _ t j => t.val * BLOCK_SIZE + j.val < N
Referenced model/specification definitions: Python source |
|---|
lightning_attentionlightning_attention_output_summary_generalMathematical execution | Configured model / stage Forward and backward producer/step statements over the modeled block/layout configurations; each conjunct has its own kernel and hypotheses. | Statement and preconditionsspecification lightning_attention_output_summary_general
(Q K V Out DO DQ DK DV KVPrev KTrans VTile KVOut OInter : RegionName)
(s : BlockState)
(_b h n d e BLOCK NUM_BLOCK BLOCK_MODEL m : Nat)
-- block `m` lies inside the sequence; `NUM_BLOCK = cdiv(n, BLOCK)`
-- (`lightning_attention.py:378`) so only a partial tail block can violate it
-- the channel block fits the value width; the launch sets
-- `BLOCK_MODEL = min(next_power_of_2(e), 32)` (`lightning_attention.py:380`)
(hBM : BLOCK_MODEL ≤ e)
(hPrev : ∀ idx : TileIndex [d, BLOCK_MODEL],
s.readMem KVPrev (kvOffset BLOCK_MODEL idx)
= kvClosed s K V n d e BLOCK BLOCK_MODEL m idx.1.val idx.2.1.val)
(hK : ∀ (idx : TileIndex [d, BLOCK_MODEL]) (j : Fin BLOCK),
s.readMem KTrans (idx.1.val * BLOCK + j.val)
= fwdKVal s K n d idx.1.val (m * BLOCK + j.val))
(hV : ∀ (idx : TileIndex [d, BLOCK_MODEL]) (j : Fin BLOCK),
s.readMem VTile (j.val * BLOCK_MODEL + idx.2.1.val)
= fwdVVal s V n e BLOCK_MODEL idx.2.1.val (m * BLOCK + j.val)) :
-- (1) surface lowering of forward + both backward kernels (symbolic dims)
(∃ alg, (lightning_attention_forward_surface Q K V Out _b h n d e BLOCK
NUM_BLOCK BLOCK_MODEL).toAlgorithm? = Except.ok alg) ∧
(∃ alg, (lightning_attention_bwd_intra_surface Q K V DO DQ DK DV
_b h n d e BLOCK NUM_BLOCK BLOCK NUM_BLOCK).toAlgorithm? = Except.ok alg) ∧
(∃ alg, (lightning_attention_bwd_inter_surface Q K V DO DQ DK DV
_b h n d e BLOCK NUM_BLOCK BLOCK NUM_BLOCK).toAlgorithm? = Except.ok alg) ∧
-- (2) kv carry-fold body realizes kvStepSpec ...
(ComputeCorrect.Realizes_without_Rounding
(kernel := lightning_attention_forward_kv_step_slice KVPrev KTrans VTile KVOut
d BLOCK BLOCK_MODEL)
(initialState := s)
(write := fun idx : TileIndex [d, BLOCK_MODEL] =>
some (KVOut, kvOffset BLOCK_MODEL idx))
(expected := fun idx : TileIndex [d, BLOCK_MODEL] =>
kvStepSpec s KVPrev KTrans VTile d BLOCK BLOCK_MODEL idx)) ∧
-- ... and that spec is the genuine closed-form kvClosed (m+1) under the carry invariant
(∀ idx : TileIndex [d, BLOCK_MODEL],
kvStepSpec s KVPrev KTrans VTile d BLOCK BLOCK_MODEL idx
= kvClosed s K V n d e BLOCK BLOCK_MODEL (m + 1) idx.1.val idx.2.1.val) ∧
-- (3) o_inter producer realizes its genuine spec oInterDotSpec
(ComputeCorrect.Realizes_without_Rounding
(kernel := lightning_attention_forward_o_inter_dot_slice Q KVPrev OInter
BLOCK d BLOCK_MODEL)
(initialState := s)
(write := fun idx : TileIndex [BLOCK, BLOCK_MODEL] =>
some (OInter, oInterOffset BLOCK_MODEL idx))
(expected := fun idx : TileIndex [BLOCK, BLOCK_MODEL] =>
oInterDotSpec s Q KVPrev BLOCK d BLOCK_MODEL idx)) ∧
-- (4) ★ the WHOLE loop body — the causal-masked `o_intra`, `o_inter`, their
-- sum, and the masked `tl.store` into the launched `Out` — realizes the
-- genuine causal linear-attention output row
(ComputeCorrect.Realizes_without_Rounding
(kernel := lightning_attention_forward_body_slice Q K V KVPrev Out
m n d e BLOCK BLOCK_MODEL)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK, BLOCK_MODEL] => m * BLOCK + idx.1.val < n)
(fun idx => (Out, bodyOutOffset s m n e BLOCK BLOCK_MODEL idx)))
(expected := fun idx : TileIndex [BLOCK, BLOCK_MODEL] =>
fwdOutClosed s Q K V n d e BLOCK_MODEL
(m * BLOCK + idx.1.val) idx.2.1.val))
Referenced model/specification definitions: Python source |
|---|
lightning_attentionlightning_attention_bwd_grad_store_io_correctnessMathematical execution | Precomputed-input slice Stores the supplied GradPre tile; this IO kernel does not compute attention gradients. | Statement and preconditionsspecification lightning_attention_bwd_grad_store_io_correctness
(GradPre Out : RegionName) (n width BLOCK WIDTH : Nat)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK, WIDTH] =>
p₀ * n * width + (p₁ * BLOCK + idx.1.val) * width + idx.2.1.val)) :
gradStoreIO GradPre Out n width BLOCK WIDTH
⊨ fun _p₀ _p₁ xs idx => xs idx
IO definition: gradStoreIOdef gradStoreIO (GradPre Out : RegionName) (n width BLOCK WIDTH : Nat) :
Masked3DTileKernelIO₁ where
kernel := lightning_attention_bwd_grad_store_slice GradPre Out n width BLOCK
WIDTH
inp := GradPre
out := Out
shape := [BLOCK, WIDTH]
read := fun p₀ p₁ _p₂ idx =>
p₀ * n * width + (p₁ * BLOCK + idx.1.val) * width + idx.2.1.val
write := fun p₀ p₁ _p₂ idx =>
p₀ * n * width + (p₁ * BLOCK + idx.1.val) * width + idx.2.1.val
mask := fun _p₀ p₁ _p₂ idx => p₁ * BLOCK + idx.1.val < n
Referenced model/specification definitions: Python source |
|---|
lightning_attentionlightning_attention_bwd_grad_store_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding contract for the supplied-gradient store slice. | Statement and preconditionsspecification lightning_attention_bwd_grad_store_io_correctnessR
(R : RoundingModel) (GradPre Out : RegionName) (n width BLOCK WIDTH : Nat)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK, WIDTH] =>
p₀ * n * width + (p₁ * BLOCK + idx.1.val) * width + idx.2.1.val)) :
gradStoreIO GradPre Out n width BLOCK WIDTH
⊨[R, FloatDType.real] fun _p₀ _p₁ xs idx => xs idx
IO definition: gradStoreIOdef gradStoreIO (GradPre Out : RegionName) (n width BLOCK WIDTH : Nat) :
Masked3DTileKernelIO₁ where
kernel := lightning_attention_bwd_grad_store_slice GradPre Out n width BLOCK
WIDTH
inp := GradPre
out := Out
shape := [BLOCK, WIDTH]
read := fun p₀ p₁ _p₂ idx =>
p₀ * n * width + (p₁ * BLOCK + idx.1.val) * width + idx.2.1.val
write := fun p₀ p₁ _p₂ idx =>
p₀ * n * width + (p₁ * BLOCK + idx.1.val) * width + idx.2.1.val
mask := fun _p₀ p₁ _p₂ idx => p₁ * BLOCK + idx.1.val < n
Referenced model/specification definitions: Python source |
|---|
llama_ff_tritonff_llama_closed_form_correctMathematical execution | Configured model / stage Selected fused feed-forward matmul/RMS/activation value target with the displayed reduction/stride configuration and mathematical fp16 cast; no frame is stated by Realizes alone. | Statement and preconditionsspecification ff_llama_closed_form_correct
(A W1 W3 OUT RMS : RegionName) (s : BlockState)
(M N sam sak sw1k sw1n sw3k sw3n soutm soutn srms BM BN BK numKBlocks : Nat) (K : Nat)
(hK : K = BK * numKBlocks) (EPS : ℝ)
(hcn : soutn = 1) (hbnle : BN ≤ soutm)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := ff_llama_surface A W1 W3 OUT RMS M N K sam sak sw1k sw1n sw3k sw3n
soutm soutn srms BM BN BK numKBlocks EPS)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BM BN)
(fun idx => (OUT, outOffset s N BM BN soutm soutn idx)))
(expected := fun idx : TileIndex [BM, BN] =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (ffSpec s A W1 W3 RMS M N BM BN sam sak sw1k sw1n sw3k sw3n srms
BK numKBlocks K EPS idx.1 idx.2.1))))
Referenced model/specification definitions: Python source |
|---|
log_softmaxlog_softmax_backward_kernel_correctnessMathematical execution | Configured model / stage Separate per-tile log-softmax and backward contracts with explicit masks, layout injectivity, and mathematical exp/log semantics. Wrapper reductions and concrete floating behavior are external. | Statement and preconditionsspecification log_softmax_backward_kernel_correctness
(out_ptr out_grad_ptr in_grad_ptr : RegionName)
(M N K BLOCK_M BLOCK_N : Nat)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun j : Fin (BLOCK_M * BLOCK_N) =>
outOffset p₀ p₁ N K BLOCK_M (laneIdx BLOCK_M BLOCK_N j))) :
Masked2DKernelIO₂.Implements
(logSoftmaxBackwardIO out_ptr out_grad_ptr in_grad_ptr M N K BLOCK_M BLOCK_N)
(fun p₀ _ xs ys j =>
logSoftmaxBackwardOf BLOCK_M BLOCK_N
(logSoftmaxBackwardPureOutTile p₀ M N BLOCK_M BLOCK_N xs)
(logSoftmaxBackwardPureGradTile p₀ M N BLOCK_M BLOCK_N ys)
(laneIdx BLOCK_M BLOCK_N j))
IO definition: logSoftmaxBackwardIOdef logSoftmaxBackwardIO (out_ptr out_grad_ptr in_grad_ptr : RegionName)
(M N K BLOCK_M BLOCK_N : Nat) : Masked2DKernelIO₂ where
kernel := log_softmax_backward_kernel out_ptr out_grad_ptr in_grad_ptr M N K BLOCK_M BLOCK_N
in1 := out_ptr
in2 := out_grad_ptr
out := in_grad_ptr
B := BLOCK_M * BLOCK_N
read1 := fun p₀ p₁ j => outOffset p₀ p₁ N K BLOCK_M (laneIdx BLOCK_M BLOCK_N j)
read2 := fun p₀ p₁ j => outOffset p₀ p₁ N K BLOCK_M (laneIdx BLOCK_M BLOCK_N j)
write := fun p₀ p₁ j => outOffset p₀ p₁ N K BLOCK_M (laneIdx BLOCK_M BLOCK_N j)
mask := fun p₀ _ j => active p₀ M N BLOCK_M (laneIdx BLOCK_M BLOCK_N j)
Referenced model/specification definitions: Python source |
|---|
log_softmaxlog_softmax_kernel_correctnessMathematical execution | Configured model / stage Separate per-tile log-softmax and backward contracts with explicit masks, layout injectivity, and mathematical exp/log semantics. Wrapper reductions and concrete floating behavior are external. | Statement and preconditionsspecification log_softmax_kernel_correctness
(output_ptr input_ptr : RegionName)
(M N K BLOCK_M BLOCK_N : Nat) (hBN : 0 < BLOCK_N)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun j : Fin (BLOCK_M * BLOCK_N) =>
outOffset p₀ p₁ N K BLOCK_M (laneIdx BLOCK_M BLOCK_N j))) :
logSoftmaxIO output_ptr input_ptr M N K BLOCK_M BLOCK_N
⊨ fun p₀ _ xs j =>
logSoftmaxOf BLOCK_M BLOCK_N
(logSoftmaxPureTile p₀ M N BLOCK_M BLOCK_N xs)
(laneIdx BLOCK_M BLOCK_N j)
IO definition: logSoftmaxIOdef logSoftmaxIO (output_ptr input_ptr : RegionName)
(M N K BLOCK_M BLOCK_N : Nat) : Masked2DKernelIO₁ where
kernel := log_softmax_kernel output_ptr input_ptr M N K BLOCK_M BLOCK_N
inp := input_ptr
out := output_ptr
B := BLOCK_M * BLOCK_N
read := fun p₀ p₁ j => outOffset p₀ p₁ N K BLOCK_M (laneIdx BLOCK_M BLOCK_N j)
write := fun p₀ p₁ j => outOffset p₀ p₁ N K BLOCK_M (laneIdx BLOCK_M BLOCK_N j)
mask := fun p₀ _ j => active p₀ M N BLOCK_M (laneIdx BLOCK_M BLOCK_N j)
Referenced model/specification definitions: Python source |
|---|
logsumexp_fwdlogsumexp_fwd_kernel_correctnessMathematical execution | Configured model / stage Per-program active-block log-sum-exp, optionally scaled. A whole-tensor reduction across program blocks is not part of this output contract. | Statement and preconditionsspecification logsumexp_fwd_kernel_correctness
(x z : RegionName) (D B : Nat) (HAS_SCALE : Bool) (scale : ℝ)
(hB : 0 < B) :
logsumexpIO x z D B HAS_SCALE scale ⊨
fun _ i_d xs _ =>
if i_d * B < D then blockLSE_local D B i_d HAS_SCALE scale xs else 0
IO definition: logsumexpIOdef logsumexpIO (x z : RegionName) (D B : Nat) (HAS_SCALE : Bool) (scale : ℝ) :
Masked2DKernelIO₁ where
kernel := logsumexp_fwd_kernel x z D B HAS_SCALE scale
inp := x
out := z
B := B
read := fun i_n i_d j => i_n * D + (i_d * B + j.val)
write := fun i_n i_d _ => i_n * ((D + B - 1) / B) + i_d
mask := fun _ i_d j => i_d * B + j.val < D
writeMask := fun _ _ j => j.val = 0
Referenced model/specification definitions: Python source |
|---|
lora_expand_gemvgemv_full_output_summaryMathematical execution | Configured model / stage Selected LoRA expansion stream and output layout with supplied adapter metadata. The exact K/BLOCK_K and block-size hypotheses constrain the original-kernel configuration. | Statement and preconditionsspecification gemv_full_output_summary
(input_ptr lora_ptr out_ptr : RegionName) (lora_indices : Region .nat)
(K split_n_length xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K : Nat)
(s : BlockState) (hBN : 0 < BLOCK_N) (hKB : K ≤ BLOCK_K) (hol : out_ptr ≠ lora_ptr)
(hundef : ∀ rg o, s.undef rg o = 0)
(hcn : 0 < cn_stride) :
(∃ alg, (bgmv_loop_surface input_ptr lora_ptr out_ptr lora_indices K
split_n_length xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := bgmv_loop_surface input_ptr lora_ptr out_ptr lora_indices K
split_n_length xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun _ : Fin split_n_length => True)
(fun m => (out_ptr, outOffG s split_n_length cm_stride cn_stride m.val)))
(expected := fun m : Fin split_n_length =>
gemvSpec s input_ptr lora_ptr lora_indices K split_n_length
xm_stride xk_stride l0_stride lora_k_stride lora_n_stride m.val)
Referenced model/specification definitions: Python source |
|---|
lora_expand_gemvgemv_full_io_correctnessAbstract cast/store rounding | Configured model / stage Selected LoRA expansion stream and output layout with supplied adapter metadata. The exact K/BLOCK_K and block-size hypotheses constrain the original-kernel configuration. | Statement and preconditionsspecification gemv_full_io_correctness (R : RoundingModel)
(input_ptr lora_ptr out_ptr : RegionName) (lora_indices : Region .nat)
(K split_n_length xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K : Nat)
(hBN : 0 < BLOCK_N) (hKB : K ≤ BLOCK_K) (hsnl : 0 < split_n_length)
(hol : out_ptr ≠ lora_ptr) (hcn : 0 < cn_stride) :
loraExpandGemvIO input_ptr lora_ptr out_ptr lora_indices K split_n_length
xm_stride xk_stride l0_stride lora_k_stride lora_n_stride cm_stride
cn_stride BLOCK_N BLOCK_K ⊨[R]
fun _ _ _ _ xs ys t j =>
∑ k : Fin BLOCK_K, if k.val < K then
xs t k * ys t (bLane BLOCK_N BLOCK_K j k) else 0
IO definition: loraExpandGemvIOdef loraExpandGemvIO (input_ptr lora_ptr out_ptr : RegionName)
(lora_indices : Region .nat)
(K split_n_length xm_stride xk_stride l0_stride lora_k_stride
lora_n_stride cm_stride cn_stride BLOCK_N BLOCK_K : Nat) :
StreamMetaEmitMasked3DKernelIO₂ where
kernel := bgmv_loop_surface input_ptr lora_ptr out_ptr lora_indices K
split_n_length xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride BLOCK_N BLOCK_K
inp1 := input_ptr
inp2 := lora_ptr
out := out_ptr
nMeta := 1
sty := fun _ => ChanTy.nat
mbuf := fun _ => lora_indices.cast
mwin := fun _ _ pid₁ _ => pid₁
T := bgmvNumSteps split_n_length BLOCK_N
B1 := BLOCK_K
B2 := BLOCK_N * BLOCK_K
C := BLOCK_N
read1 := fun _ pid₁ _ _ _ k => pid₁ * xm_stride + k.val * xk_stride
read2 := fun pid₀ _ _ m t l =>
l0_stride * m (⟨0, by omega⟩ : Fin 1)
+ pid₀ * split_n_length * lora_k_stride
+ (t.val * BLOCK_N + l.val / BLOCK_K) * lora_k_stride
+ (l.val % BLOCK_K) * lora_n_stride
write := fun pid₀ pid₁ _ _ t j =>
pid₁ * cm_stride + pid₀ * split_n_length + (t.val * BLOCK_N + j.val) * cn_stride
mask1 := fun _ _ _ _ _ k => k.val < K
mask2 := fun _ _ _ _ t l =>
t.val * BLOCK_N + l.val / BLOCK_K < split_n_length ∧ l.val % BLOCK_K < K
writeMask := fun _ _ _ _ t j => t.val * BLOCK_N + j.val < split_n_length
Referenced model/specification definitions: Python source |
|---|
masked_add_cudamasked_add_kernel_correctnessMathematical execution | Configured model / stage Masked in-place parameter update from supplied gradients and mask values. Only active-lane writes are claimed; mask construction and host scheduling are external. | Statement and preconditionsspecification masked_add_kernel_correctness
(grad_ptr p_ptr p_mask_ptr : RegionName)
(n_elements : Nat) (alpha : ℝ) (BLOCK_SIZE : Nat) :
maskedAddIO grad_ptr p_ptr p_mask_ptr n_elements alpha BLOCK_SIZE
⊨ fun _ _ _ xs ys j => ys j + xs j * alpha
IO definition: maskedAddIOdef maskedAddIO (grad_ptr p_ptr p_mask_ptr : RegionName)
(n_elements : Nat) (alpha : ℝ) (BLOCK_SIZE : Nat) :
BoolMasked2DKernelIO₂ where
kernel := masked_add_kernel grad_ptr p_ptr p_mask_ptr n_elements alpha
BLOCK_SIZE
in1 := p_ptr
in2 := grad_ptr
mbuf := p_mask_ptr
out := grad_ptr
B := BLOCK_SIZE
read1 := fun pid₀ _ j => pid₀ * BLOCK_SIZE + j.val
read2 := fun pid₀ _ j => pid₀ * BLOCK_SIZE + j.val
readm := fun pid₀ _ j => pid₀ * BLOCK_SIZE + j.val
write := fun pid₀ _ j => pid₀ * BLOCK_SIZE + j.val
mask := fun pid₀ _ j => pid₀ * BLOCK_SIZE + j.val < n_elements
writeMask := fun pid₀ _ bs j =>
pid₀ * BLOCK_SIZE + j.val < n_elements ∧ bs j = Bool.false
Referenced model/specification definitions: Python source |
|---|
masked_selectmasked_select_kernel_correctnessMathematical execution | Configured model / stage Selection-copy stage consumes supplied selection mask and prefix-sum metadata. It does not compute or validate a prefix sum beyond the IO layout/bounds premises. | Statement and preconditionsspecification masked_select_kernel_correctness
(inp_ptr select_mask_ptr prefix_sum_ptr out_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) :
maskedSelectIO inp_ptr select_mask_ptr prefix_sum_ptr out_ptr n_elements
BLOCK_SIZE ⊨ fun _ _ _ _ xs j => xs j
IO definition: maskedSelectIOdef maskedSelectIO (inp_ptr select_mask_ptr prefix_sum_ptr out_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) : BoolScatterMasked2DKernelIO₁ where
kernel := masked_select_kernel inp_ptr select_mask_ptr prefix_sum_ptr out_ptr
n_elements BLOCK_SIZE
inp := inp_ptr
mbuf := select_mask_ptr
idxbuf := prefix_sum_ptr
out := out_ptr
B := BLOCK_SIZE
read := fun pid₀ _ j => pid₀ * BLOCK_SIZE + j.val
readm := fun pid₀ _ j => pid₀ * BLOCK_SIZE + j.val
readx := fun pid₀ _ j => pid₀ * BLOCK_SIZE + j.val
write := fun _ _ ids j => ids j - 1
mask := fun pid₀ _ j => pid₀ * BLOCK_SIZE + j.val < n_elements
writeMask := fun pid₀ _ bs _ j =>
pid₀ * BLOCK_SIZE + j.val < n_elements ∧ bs j = Bool.true
Referenced model/specification definitions: Python source |
|---|
matmul_dequant_int4matmul_dequant_int4_exec_genuineMathematical execution | Configured model / stage The registered headline is the dequantize output stage, despite the file containing a matmul Python function. Its exec-existential omits a frame and uses modeled integer decode semantics. | Statement and preconditionsspecification matmul_dequant_int4_exec_genuine
(b_ptr : Region .nat) (b_scale_ptr : RegionName) (b_zp_ptr : Region .nat)
(fpb_ptr : RegionName)
(K N group_size stride_bk stride_bn stride_bsk stride_bsn
stride_bzpk stride_bzpn stride_fpbk stride_fpbn : Nat)
(BK BN : Nat) (s : BlockState)
(hInj : Function.Injective
(fun idx : TileIndex [BK, BN] =>
fpbAddr stride_fpbk stride_fpbn BK BN (s.pids 0) (s.pids 1) idx)) :
∃ sF, exec ((matmul_dequant_int4_surface b_ptr b_scale_ptr b_zp_ptr fpb_ptr
K N group_size stride_bk stride_bn stride_bsk stride_bsn stride_bzpk
stride_bzpn stride_fpbk stride_fpbn BK BN).toAlgKernel) s = some sF
∧ ∀ idx : TileIndex [BK, BN],
(s.pids 0 * BK + idx.1.val < K ∧ s.pids 1 * BN + idx.2.1.val < N) →
sF.readMem fpb_ptr
(fpbAddr stride_fpbk stride_fpbn BK BN (s.pids 0) (s.pids 1) idx)
= dequantSpec s b_scale_ptr b_ptr b_zp_ptr group_size stride_bk
stride_bn stride_bsk stride_bsn stride_bzpk stride_bzpn
(s.pids 0 * BK + idx.1.val) (s.pids 1 * BN + idx.2.1.val)
Referenced model/specification definitions: Python source |
|---|
matmul_dequantizematmul_dequantize_matmul4_exec_genuineMathematical execution | Configured model / stage Separate matmul4, int4 matmul, and dequantize stage readbacks with their explicit grouping/stride assumptions. Exec-existentials omit frame; no wrapper pipeline or machine-overflow theorem. | Statement and preconditionsspecification matmul_dequantize_matmul4_exec_genuine
(a_ptr c_ptr scales_ptr : RegionName) (b_ptr zeros_ptr : Region .nat)
(M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
stride_scales_g stride_scales_n stride_zeros_g stride_zeros_n
groupsize : Nat) (NO_GROUPS : Bool) (BM BN BK GM : Nat) (s : BlockState)
(hInj : Function.Injective
(fun i : TileIndex [BM, BN] => cAddr stride_cm stride_cn BM BN
(pidM s M N BM BN GM) (pidN s M N BM BN GM) i)) :
∃ sF, exec (matmul_dequantize_matmul4_surface a_ptr c_ptr scales_ptr b_ptr zeros_ptr
M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
stride_scales_g stride_scales_n stride_zeros_g stride_zeros_n
groupsize NO_GROUPS BM BN BK GM).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BM, BN],
(pidM s M N BM BN GM * BM + idx.1.val < M
∧ pidN s M N BM BN GM * BN + idx.2.1.val < N) →
sF.readMem c_ptr (cAddr stride_cm stride_cn BM BN (pidM s M N BM BN GM)
(pidN s M N BM BN GM) idx)
= accSpec s a_ptr scales_ptr b_ptr zeros_ptr NO_GROUPS M K groupsize
stride_am stride_ak stride_bk stride_bn stride_scales_g
stride_scales_n stride_zeros_g stride_zeros_n BM BN BK
(pidM s M N BM BN GM) (pidN s M N BM BN GM)
idx.1.val idx.2.1.val
Referenced model/specification definitions: Python source |
|---|
matmul_dequantizematmul_dequantize_matmul_exec_genuineMathematical execution | Configured model / stage Separate matmul4, int4 matmul, and dequantize stage readbacks with their explicit grouping/stride assumptions. Exec-existentials omit frame; no wrapper pipeline or machine-overflow theorem. | Statement and preconditionsspecification matmul_dequantize_matmul_exec_genuine
(a_ptr c_ptr bs_ptr : RegionName) (b_ptr bzp_ptr : Region .nat)
(M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
stride_bsk stride_bsn stride_bzpk stride_bzpn group_size : Nat)
(BM BN BK GM numKBlocks : Nat) (s : BlockState)
(hK : K = BK * numKBlocks)
(hBK8 : BK % 8 = 0)
(hpid1 : s.pids 1 = 0)
(hInj : Function.Injective
(fun i : TileIndex [BM, BN] => cAddr stride_cm stride_cn BM BN
(pidM s M N BM BN GM) (pidN s M N BM BN GM) i)) :
∃ sF, exec (matmul_dequantize_matmul_surface a_ptr c_ptr bs_ptr b_ptr bzp_ptr
M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
stride_bsk stride_bsn stride_bzpk stride_bzpn group_size
BM BN BK GM numKBlocks).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BM, BN],
(pidM s M N BM BN GM * BM + idx.1.val < M
∧ pidN s M N BM BN GM * BN + idx.2.1.val < N) →
sF.mem c_ptr (cAddr stride_cm stride_cn BM BN (pidM s M N BM BN GM)
(pidN s M N BM BN GM) idx)
= MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (i4AccSpec s a_ptr bs_ptr b_ptr bzp_ptr group_size
stride_am stride_ak stride_bk stride_bn stride_bsk stride_bsn
stride_bzpk stride_bzpn BK numKBlocks
(pidM s M N BM BN GM * BM + idx.1.val)
(pidN s M N BM BN GM * BN + idx.2.1.val))))
Referenced model/specification definitions: Python source |
|---|
matmul_dequantizematmul_dequantize_dequantize_exec_genuineMathematical execution | Configured model / stage Separate matmul4, int4 matmul, and dequantize stage readbacks with their explicit grouping/stride assumptions. Exec-existentials omit frame; no wrapper pipeline or machine-overflow theorem. | Statement and preconditionsspecification matmul_dequantize_dequantize_exec_genuine
(b_ptr : Region .nat) (b_scale_ptr : RegionName) (b_zp_ptr : Region .nat)
(fpb_ptr : RegionName)
(K N group_size stride_bk stride_bn stride_bsk stride_bsn
stride_bzpk stride_bzpn stride_fpbk stride_fpbn : Nat)
(BK BN : Nat) (s : BlockState)
(hInj : Function.Injective
(fun idx : TileIndex [BK, BN] =>
fpbAddr stride_fpbk stride_fpbn BK BN (s.pids 0) (s.pids 1) idx)) :
∃ sF, exec ((matmul_dequantize_dequantize_surface b_ptr b_scale_ptr b_zp_ptr fpb_ptr
K N group_size stride_bk stride_bn stride_bsk stride_bsn stride_bzpk
stride_bzpn stride_fpbk stride_fpbn BK BN).toAlgKernel) s = some sF
∧ ∀ idx : TileIndex [BK, BN],
(s.pids 0 * BK + idx.1.val < K ∧ s.pids 1 * BN + idx.2.1.val < N) →
sF.readMem fpb_ptr
(fpbAddr stride_fpbk stride_fpbn BK BN (s.pids 0) (s.pids 1) idx)
= dequantSpec s b_scale_ptr b_ptr b_zp_ptr group_size stride_bk
stride_bn stride_bsk stride_bsn stride_bzpk stride_bzpn
(s.pids 0 * BK + idx.1.val) (s.pids 1 * BN + idx.2.1.val)
Referenced model/specification definitions: Python source |
|---|
matmul_dequantize_int4matmul_dequantize_int4_exec_genuineMathematical execution | Configured model / stage Matmul4 stage at the stated group/tile/layout specialization with modeled packed-integer decoding. Its exec-existential does not state a memory frame or hardware overflow behavior. | Statement and preconditionsspecification matmul_dequantize_int4_exec_genuine
(a_ptr c_ptr scales_ptr : RegionName) (b_ptr zeros_ptr : Region .nat)
(M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
stride_scales_g stride_scales_n stride_zeros_g stride_zeros_n
groupsize : Nat) (NO_GROUPS : Bool) (BM BN BK GM : Nat) (s : BlockState)
(hInj : Function.Injective
(fun i : TileIndex [BM, BN] => cAddr stride_cm stride_cn BM BN
(pidM s M N BM BN GM) (pidN s M N BM BN GM) i)) :
∃ sF, exec (matmul_dequantize_int4_surface a_ptr c_ptr scales_ptr b_ptr zeros_ptr
M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
stride_scales_g stride_scales_n stride_zeros_g stride_zeros_n
groupsize NO_GROUPS BM BN BK GM).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BM, BN],
(pidM s M N BM BN GM * BM + idx.1.val < M
∧ pidN s M N BM BN GM * BN + idx.2.1.val < N) →
sF.readMem c_ptr (cAddr stride_cm stride_cn BM BN (pidM s M N BM BN GM)
(pidN s M N BM BN GM) idx)
= accSpec s a_ptr scales_ptr b_ptr zeros_ptr NO_GROUPS M K groupsize
stride_am stride_ak stride_bk stride_bn stride_scales_g
stride_scales_n stride_zeros_g stride_zeros_n BM BN BK
(pidM s M N BM BN GM) (pidN s M N BM BN GM)
idx.1.val idx.2.1.val
Referenced model/specification definitions: Python source |
|---|
matmul_kernelmatmul_kernel_closed_form_correctMathematical execution | Configured model / stage Configured contiguous matmul surface; the IO theorem explicitly requires BN<=4096 and names the fixed-layout kernel. This is not arbitrary Python stride/autotuning coverage. | Statement and preconditionsspecification matmul_kernel_closed_form_correct
(C A B : RegionName) (s : BlockState)
(BM BN BLOCK_K numKBlocks : Nat)
(hBN : BN ≤ 4096)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := matmul_kernel_surface C A B BM BN BLOCK_K numKBlocks)
(initialState := s)
(write := fun idx : TileIndex [BM, BN] => some (C, cOffset s BM BN idx))
(expected := fun idx : TileIndex [BM, BN] =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (matmulSpec s A B BM BN BLOCK_K numKBlocks idx.1 idx.2.1))))
Referenced model/specification definitions: Python source |
|---|
matmul_kernelmatmul_kernel_io_correctnessAbstract cast/store rounding | Configured model / stage Configured contiguous matmul surface; the IO theorem explicitly requires BN<=4096 and names the fixed-layout kernel. This is not arbitrary Python stride/autotuning coverage. | Statement and preconditionsspecification matmul_kernel_io_correctness (R : RoundingModel)
(C A B : RegionName) (BM BN BK T : Nat) (hBN : BN ≤ 4096) :
matmulKernelIO C A B BM BN BK T ⊨[R] fun _ _ xs ys l =>
∑ t : Fin T, ∑ e : Fin BK, xs t (aLane BM BN BK l e) * ys t (bLane BM BN BK l e)
IO definition: matmulKernelIOdef matmulKernelIO (C A B : RegionName) (BM BN BK T : Nat) : StreamMasked2DKernelIO₂ where
kernel := matmul_kernel_surface C A B BM BN BK T
inp1 := A
inp2 := B
out := C
T := T
B1 := BM * BK
B2 := BK * BN
C := BM * BN
outDType := .fp16
read1 := fun p₀ _ t l => ((p₀ * BM + l.val / BK) % 4096) * 4096 + (t.val * BK + l.val % BK)
read2 := fun _ p₁ t l => (t.val * BK + l.val / BN) * 4096 + (p₁ * BN + l.val % BN) % 4096
write := fun p₀ p₁ l => 4096 * (p₀ * BM + l.val / BN) + 1 * (p₁ * BN + l.val % BN)
mask1 := fun _ _ _ _ => True
mask2 := fun _ _ _ _ => True
writeMask := fun _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
matmul_leakyrelumatmul_leakyrelu_closed_form_correctMathematical execution | Configured model / stage Matmul plus leaky-ReLU for K=BK*numKBlocks and the stated output-layout bounds. Host autotuning, alternate layouts, and concrete floating behavior are external. | Statement and preconditionsspecification matmul_leakyrelu_closed_form_correct
(A B C : RegionName) (s : BlockState)
(M N SAM SAK SBK SBN SCM SCN BM BN BK GROUP numKBlocks : Nat) (hBK : 0 < BK)
(hscn : SCN = 1) (hbnle : BN ≤ SCM)
(hmlt : ∀ i : Fin BM, rowIndex (pidM (s.pids 0) M N BM BN GROUP) BM i < M)
(hnlt : ∀ j : Fin BN, colIndex (pidN (s.pids 0) M N BM BN GROUP) BN j < N)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := matmul_leaky_relu_surface A B C M N (BK * numKBlocks) SAM SAK SBK SBN SCM SCN BM BN BK GROUP)
(initialState := s)
(write := fun idx : TileIndex [BM, BN] =>
some (C, cOffset s (pidM (s.pids 0) M N BM BN GROUP) (pidN (s.pids 0) M N BM BN GROUP) BM BN SCM SCN idx))
(expected := fun idx : TileIndex [BM, BN] =>
outputCell s A B (pidM (s.pids 0) M N BM BN GROUP) (pidN (s.pids 0) M N BM BN GROUP)
BM BN M N SAM SAK SBK SBN BK numKBlocks idx)
Referenced model/specification definitions: Python source |
|---|
matmul_leakyrelumatmul_leakyrelu_io_correctnessAbstract cast/store rounding | Configured model / stage Matmul plus leaky-ReLU for K=BK*numKBlocks and the stated output-layout bounds. Host autotuning, alternate layouts, and concrete floating behavior are external. | Statement and preconditionsspecification matmul_leakyrelu_io_correctness (R : RoundingModel)
(A B C : RegionName)
(M N SAM SAK SBK SBN SCM SCN BM BN BK GROUP numKBlocks : Nat)
(hBK : 0 < BK) (hscn : SCN = 1) (hbnle : BN ≤ SCM) :
matmulLeakyreluIO A B C M N SAM SAK SBK SBN SCM SCN BM BN BK GROUP numKBlocks
⊨[R] fun _ _ xs ys l =>
leakyrelu (∑ t : Fin numKBlocks, ∑ e : Fin BK,
xs t (aLane BM BN BK l e) * ys t (bLane BM BN BK l e))
IO definition: matmulLeakyreluIOdef matmulLeakyreluIO (A B C : RegionName)
(M N SAM SAK SBK SBN SCM SCN BM BN BK GROUP numKBlocks : Nat) :
StreamMasked2DKernelIO₂ where
kernel := matmul_leaky_relu_surface A B C M N (BK * numKBlocks) SAM SAK SBK SBN SCM SCN
BM BN BK GROUP
inp1 := A
inp2 := B
out := C
T := numKBlocks
B1 := BM * BK
B2 := BK * BN
C := BM * BN
outDType := .fp16
read1 := fun p₀ _ t l =>
(pidM p₀ M N BM BN GROUP * BM + l.val / BK) % M * SAM + (t.val * BK + l.val % BK) * SAK
read2 := fun p₀ _ t l =>
(t.val * BK + l.val / BN) * SBK + (pidN p₀ M N BM BN GROUP * BN + l.val % BN) % N * SBN
write := fun p₀ _ l =>
SCM * (pidM p₀ M N BM BN GROUP * BM + l.val / BN)
+ SCN * (pidN p₀ M N BM BN GROUP * BN + l.val % BN)
mask1 := fun _ _ t l => t.val * BK + l.val % BK < BK * numKBlocks
mask2 := fun _ _ t l => t.val * BK + l.val / BN < BK * numKBlocks
writeMask := fun p₀ _ l =>
pidM p₀ M N BM BN GROUP * BM + l.val / BN < M ∧
pidN p₀ M N BM BN GROUP * BN + l.val % BN < N
Referenced model/specification definitions: Python source |
|---|
matmul_leakyrelu_fp8matmul_leakyrelu_fp8_closed_form_correctMathematical execution | Configured model / stage Matmul plus leaky-ReLU on the declared mathematical carrier and output layout; the fp8 filename does not establish concrete fp8 encoding or rounding behavior. | Statement and preconditionsspecification matmul_leakyrelu_fp8_closed_form_correct
(A B C : RegionName) (s : BlockState)
(M N SAM SAK SBK SBN SCM SCN BM BN BK GROUP numKBlocks : Nat) (hBK : 0 < BK)
(hcn : SCN = 1) (hbnle : BN ≤ SCM)
(hmlt : ∀ i : Fin BM, rowIndex (pidM (s.pids 0) M N BM BN GROUP) BM i < M)
(hnlt : ∀ j : Fin BN, colIndex (pidN (s.pids 0) M N BM BN GROUP) BN j < N)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := matmul_leaky_relu_surface A B C M N (BK * numKBlocks) SAM SAK SBK SBN SCM SCN BM BN BK GROUP)
(initialState := s)
(write := fun idx : TileIndex [BM, BN] =>
some (C, cOffset s (pidM (s.pids 0) M N BM BN GROUP) (pidN (s.pids 0) M N BM BN GROUP) BM BN SCM SCN idx))
(expected := fun idx : TileIndex [BM, BN] =>
outputCell s A B (pidM (s.pids 0) M N BM BN GROUP) (pidN (s.pids 0) M N BM BN GROUP)
BM BN M N SAM SAK SBK SBN BK numKBlocks idx)
Referenced model/specification definitions: Python source |
|---|
matmul_leakyrelu_fp8matmul_leakyrelu_fp8_io_correctnessAbstract cast/store rounding | Configured model / stage Matmul plus leaky-ReLU on the declared mathematical carrier and output layout; the fp8 filename does not establish concrete fp8 encoding or rounding behavior. | Statement and preconditionsspecification matmul_leakyrelu_fp8_io_correctness (R : RoundingModel)
(A B C : RegionName)
(M N SAM SAK SBK SBN SCM SCN BM BN BK GROUP numKBlocks : Nat)
(hBK : 0 < BK) (hscn : SCN = 1) (hbnle : BN ≤ SCM) :
matmulLeakyreluFp8IO A B C M N SAM SAK SBK SBN SCM SCN BM BN BK GROUP numKBlocks
⊨[R] fun _ _ xs ys l =>
leakyrelu (∑ t : Fin numKBlocks, ∑ e : Fin BK,
xs t (aLane BM BN BK l e) * ys t (bLane BM BN BK l e))
IO definition: matmulLeakyreluFp8IOdef matmulLeakyreluFp8IO (A B C : RegionName)
(M N SAM SAK SBK SBN SCM SCN BM BN BK GROUP numKBlocks : Nat) :
StreamMasked2DKernelIO₂ where
kernel := matmul_leaky_relu_surface A B C M N (BK * numKBlocks) SAM SAK SBK SBN SCM SCN
BM BN BK GROUP
inp1 := A
inp2 := B
out := C
T := numKBlocks
B1 := BM * BK
B2 := BK * BN
C := BM * BN
outDType := .fp16
read1 := fun p₀ _ t l =>
(pidM p₀ M N BM BN GROUP * BM + l.val / BK) % M * SAM + (t.val * BK + l.val % BK) * SAK
read2 := fun p₀ _ t l =>
(t.val * BK + l.val / BN) * SBK + (pidN p₀ M N BM BN GROUP * BN + l.val % BN) % N * SBN
write := fun p₀ _ l =>
SCM * (pidM p₀ M N BM BN GROUP * BM + l.val / BN)
+ SCN * (pidN p₀ M N BM BN GROUP * BN + l.val % BN)
mask1 := fun _ _ t l => t.val * BK + l.val % BK < BK * numKBlocks
mask2 := fun _ _ t l => t.val * BK + l.val / BN < BK * numKBlocks
writeMask := fun p₀ _ l =>
pidM p₀ M N BM BN GROUP * BM + l.val / BN < M ∧
pidN p₀ M N BM BN GROUP * BN + l.val % BN < N
Referenced model/specification definitions: Python source |
|---|
matmul_tmamatmul_tma_f32_closed_form_correctMathematical execution | Configured model / stage One modeled matmul tile with separate exact/fp16-store contracts. Algorithm projection abstracts descriptor/TMA behavior; hardware TMA execution and synchronization are not proved. | Statement and preconditionsspecification matmul_tma_f32_closed_form_correct
(A B C : RegionName) (s : BlockState)
(M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
BLOCK_M BLOCK_N BLOCK_K : Nat)
(hcn : stride_cn = 1) (hcm : BLOCK_N ≤ stride_cm) :
ComputeCorrect.Realizes_without_Rounding
(kernel := matmul_tma_f32_surface A B C M N K stride_am stride_ak
stride_bk stride_bn stride_cm stride_cn BLOCK_M BLOCK_N BLOCK_K)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
some (C, cOffset stride_cm stride_cn idx))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
matmulSpec s A B stride_am stride_ak stride_bk stride_bn BLOCK_K
idx.1.val idx.2.1.val)
Referenced model/specification definitions: Python source |
|---|
matmul_tmamatmul_tma_f16_closed_form_correctMathematical execution | Configured model / stage One modeled matmul tile with separate exact/fp16-store contracts. Algorithm projection abstracts descriptor/TMA behavior; hardware TMA execution and synchronization are not proved. | Statement and preconditionsspecification matmul_tma_f16_closed_form_correct
(A B C : RegionName) (s : BlockState)
(M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
BLOCK_M BLOCK_N BLOCK_K : Nat)
(hcn : stride_cn = 1) (hcm : BLOCK_N ≤ stride_cm) :
ComputeCorrect.Realizes_without_Rounding
(kernel := matmul_tma_f16_surface A B C M N K stride_am stride_ak
stride_bk stride_bn stride_cm stride_cn BLOCK_M BLOCK_N BLOCK_K)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
some (C, cOffset stride_cm stride_cn idx))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (matmulSpec s A B stride_am stride_ak stride_bk stride_bn BLOCK_K
idx.1.val idx.2.1.val))))
Referenced model/specification definitions: Python source |
|---|
matmul_tmamatmul_tma_f32_io_correctnessMathematical execution | Configured model / stage One modeled matmul tile with separate exact/fp16-store contracts. Algorithm projection abstracts descriptor/TMA behavior; hardware TMA execution and synchronization are not proved. | Statement and preconditionsspecification matmul_tma_f32_io_correctness (A B C : RegionName)
(M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
BLOCK_M BLOCK_N BLOCK_K : Nat)
(hInj : Function.Injective
(cOffset (BLOCK_M := BLOCK_M) (BLOCK_N := BLOCK_N) stride_cm stride_cn)) :
matmulTmaF32IO A B C M N K stride_am stride_ak stride_bk stride_bn stride_cm
stride_cn BLOCK_M BLOCK_N BLOCK_K
⊨ fun _p₀ _p₁ xs ys idx =>
matmulSpecOf BLOCK_M BLOCK_N BLOCK_K xs ys idx
IO definition: matmulTmaF32IOdef matmulTmaF32IO (A B C : RegionName)
(M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
BLOCK_M BLOCK_N BLOCK_K : Nat) : MaskedTileShapedKernelIO₂ where
kernel := matmul_tma_f32_surface A B C M N K stride_am stride_ak stride_bk
stride_bn stride_cm stride_cn BLOCK_M BLOCK_N BLOCK_K
in1 := A
in2 := B
out := C
shape1 := [BLOCK_M, BLOCK_K]
shape2 := [BLOCK_K, BLOCK_N]
shapeOut := [BLOCK_M, BLOCK_N]
read1 := fun _p₀ _p₁ k => k.1.val * stride_am + k.2.1.val * stride_ak
read2 := fun _p₀ _p₁ k => k.1.val * stride_bk + k.2.1.val * stride_bn
write := fun _p₀ _p₁ o => cOffset stride_cm stride_cn o
mask1 := fun _p₀ _p₁ _ => True
mask2 := fun _p₀ _p₁ _ => True
writeMask := fun _p₀ _p₁ _ => True
Referenced model/specification definitions: Python source |
|---|
matmul_tmamatmul_tma_f16_io_correctnessRAbstract cast/store rounding | Configured model / stage One modeled matmul tile with separate exact/fp16-store contracts. Algorithm projection abstracts descriptor/TMA behavior; hardware TMA execution and synchronization are not proved. | Statement and preconditionsspecification matmul_tma_f16_io_correctnessR (Rm : RoundingModel)
(A B C : RegionName)
(M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
BLOCK_M BLOCK_N BLOCK_K : Nat)
(hInj : Function.Injective
(cOffset (BLOCK_M := BLOCK_M) (BLOCK_N := BLOCK_N) stride_cm stride_cn)) :
matmulTmaF16IO A B C M N K stride_am stride_ak stride_bk stride_bn stride_cm
stride_cn BLOCK_M BLOCK_N BLOCK_K
⊨[Rm, FloatDType.fp16] fun _p₀ _p₁ xs ys idx =>
matmulSpecOf BLOCK_M BLOCK_N BLOCK_K xs ys idx
IO definition: matmulTmaF16IOdef matmulTmaF16IO (A B C : RegionName)
(M N K stride_am stride_ak stride_bk stride_bn stride_cm stride_cn
BLOCK_M BLOCK_N BLOCK_K : Nat) : MaskedTileShapedKernelIO₂ where
kernel := matmul_tma_f16_surface A B C M N K stride_am stride_ak stride_bk
stride_bn stride_cm stride_cn BLOCK_M BLOCK_N BLOCK_K
in1 := A
in2 := B
out := C
shape1 := [BLOCK_M, BLOCK_K]
shape2 := [BLOCK_K, BLOCK_N]
shapeOut := [BLOCK_M, BLOCK_N]
read1 := fun _p₀ _p₁ k => k.1.val * stride_am + k.2.1.val * stride_ak
read2 := fun _p₀ _p₁ k => k.1.val * stride_bk + k.2.1.val * stride_bn
write := fun _p₀ _p₁ o => cOffset stride_cm stride_cn o
mask1 := fun _p₀ _p₁ _ => True
mask2 := fun _p₀ _p₁ _ => True
writeMask := fun _p₀ _p₁ _ => True
Referenced model/specification definitions: Python source |
|---|
matmul_triton1matmul_triton1_closed_form_correctMathematical execution | Configured model / stage Matmul at the explicit contiguous tile configuration (including fixed surface arguments and K=BLOCK_K*numKBlocks). No arbitrary Python layout or launch composition claim. | Statement and preconditionsspecification matmul_triton1_closed_form_correct
(X Y Z : RegionName) (s : BlockState)
(NS M BLOCK_K N numKBlocks : Nat) (hBK : 0 < BLOCK_K)
(hN : N ≤ NS)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := matmul_triton1_surface X Y Z 0 (BLOCK_K * numKBlocks) NS M BLOCK_K N)
(initialState := s)
(write := fun idx : TileIndex [M, N] => some (Z, zOffset s NS N M idx))
(expected := fun idx : TileIndex [M, N] =>
matmulSpec s X Y (BLOCK_K * numKBlocks) NS N M N BLOCK_K numKBlocks idx.1 idx.2.1)
Referenced model/specification definitions: Python source |
|---|
matmul_triton1matmul_triton1_io_correctnessAbstract cast/store rounding | Configured model / stage Matmul at the explicit contiguous tile configuration (including fixed surface arguments and K=BLOCK_K*numKBlocks). No arbitrary Python layout or launch composition claim. | Statement and preconditionsspecification matmul_triton1_io_correctness (R : RoundingModel)
(X Y Z : RegionName) (NS M BLOCK_K N numKBlocks : Nat)
(hBK : 0 < BLOCK_K) (hN : N ≤ NS) :
matmulTriton1IO X Y Z NS M BLOCK_K N numKBlocks ⊨[R] fun _ _ xs ys l =>
∑ t : Fin numKBlocks, ∑ e : Fin BLOCK_K,
xs t (aLane M N BLOCK_K l e) * ys t (bLane M N BLOCK_K l e)
IO definition: matmulTriton1IOdef matmulTriton1IO (X Y Z : RegionName) (NS M BLOCK_K N numKBlocks : Nat) :
StreamMasked2DKernelIO₂ where
kernel := matmul_triton1_surface X Y Z 0 (BLOCK_K * numKBlocks) NS M BLOCK_K N
inp1 := X
inp2 := Y
out := Z
T := numKBlocks
B1 := M * BLOCK_K
B2 := BLOCK_K * N
C := M * N
read1 := fun p₀ _ t l =>
(p₀ / numNBlocks NS N * M + l.val / BLOCK_K) * (BLOCK_K * numKBlocks)
+ (t.val * BLOCK_K + l.val % BLOCK_K)
read2 := fun p₀ _ t l =>
(t.val * BLOCK_K + l.val / N) * NS + (p₀ % numNBlocks NS N * N + l.val % N)
write := fun p₀ _ l =>
(p₀ / numNBlocks NS N * M + l.val / N) * NS + (p₀ % numNBlocks NS N * N + l.val % N)
mask1 := fun _ _ _ _ => True
mask2 := fun _ _ _ _ => True
writeMask := fun _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
matmul_triton2matmul_triton2_output_summary_generalMathematical execution | Configured model / stage Grouped matmul under whole reduction blocks, unit output-column stride, and output-row bounds. Host autotuning and alternate layouts are outside these contracts. | Statement and preconditionsspecification matmul_triton2_output_summary_general
(A B C : RegionName) (s : BlockState)
(M N BM BN GM SAM SAK SBK SBN SCM SCN BK numKBlocks : Nat) (hBK : 0 < BK)
(hInj : Function.Injective (cOffset s M N BM BN GM SCM SCN))
(hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (matmul_triton2_surface A B C M N (BK * numKBlocks) SAM SAK SBK SBN SCM SCN
BM BN BK GM).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := matmul_triton2_surface A B C M N (BK * numKBlocks) SAM SAK SBK SBN SCM SCN
BM BN BK GM)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BM BN GM)
(fun idx => (C, cOffset s M N BM BN GM SCM SCN idx)))
(expected := fun idx : TileIndex [BM, BN] =>
matmulSpec s A B M N BM BN GM SAM SAK SBK SBN BK numKBlocks idx.1 idx.2.1)
Referenced model/specification definitions: Python source |
|---|
matmul_triton2matmul_triton2_io_correctnessAbstract cast/store rounding | Configured model / stage Grouped matmul under whole reduction blocks, unit output-column stride, and output-row bounds. Host autotuning and alternate layouts are outside these contracts. | Statement and preconditionsspecification matmul_triton2_io_correctness (R : RoundingModel)
(A B C : RegionName)
(M N BM BN GM SAM SAK SBK SBN SCM SCN BK numKBlocks : Nat)
(hBK : 0 < BK) (hcn : SCN = 1) (hBN : BN ≤ SCM) :
matmulTriton2IO A B C M N BM BN GM SAM SAK SBK SBN SCM SCN BK numKBlocks
⊨[R] fun _ _ xs ys l =>
∑ t : Fin numKBlocks, ∑ e : Fin BK,
xs t (aLane BM BN BK l e) * ys t (bLane BM BN BK l e)
IO definition: matmulTriton2IOdef matmulTriton2IO (A B C : RegionName)
(M N BM BN GM SAM SAK SBK SBN SCM SCN BK numKBlocks : Nat) :
StreamMasked2DKernelIO₂ where
kernel := matmul_triton2_surface A B C M N (BK * numKBlocks) SAM SAK SBK SBN SCM SCN
BM BN BK GM
inp1 := A
inp2 := B
out := C
T := numKBlocks
B1 := BM * BK
B2 := BK * BN
C := BM * BN
read1 := fun p₀ _ t l =>
(pidMAt p₀ M N BM BN GM * BM + l.val / BK) * SAM + (t.val * BK + l.val % BK) * SAK
read2 := fun p₀ _ t l =>
(t.val * BK + l.val / BN) * SBK + (pidNAt p₀ M N BM BN GM * BN + l.val % BN) * SBN
write := fun p₀ _ l =>
(pidMAt p₀ M N BM BN GM * BM + l.val / BN) * SCM
+ (pidNAt p₀ M N BM BN GM * BN + l.val % BN) * SCN
mask1 := fun _ _ t l => t.val * BK + l.val % BK < BK * numKBlocks
mask2 := fun _ _ t l => t.val * BK + l.val / BN < BK * numKBlocks
writeMask := fun p₀ _ l =>
pidMAt p₀ M N BM BN GM * BM + l.val / BN < M ∧
pidNAt p₀ M N BM BN GM * BN + l.val % BN < N
Referenced model/specification definitions: Python source |
|---|
matmul_triton_autotunematmul_autotune_closed_form_correctMathematical execution | Configured model / stage Configured matmul/activation branches and output layout. The theorem quantifies over its named parameters but does not verify the Python autotuner or every autotune candidate. | Statement and preconditionsspecification matmul_autotune_closed_form_correct
(A B C : RegionName) (s : BlockState)
(M N BM BN GM sam sak sbk sbn scm scn BLOCK_K numKBlocks : Nat) (K : Nat)
(hK : K = BLOCK_K * numKBlocks) (ACTIVATION : Bool)
(hcn : scn = 1) (hbnle : BN ≤ scm)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := matmul_autotune_surface A B C M N K sam sak sbk sbn scm scn
BM BN BLOCK_K GM numKBlocks ACTIVATION)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BM BN GM)
(fun idx => (C, cOffset s M N BM BN GM scm scn idx)))
(expected := fun idx : TileIndex [BM, BN] =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (matmulSpec s A B M N BM BN GM sam sak sbk sbn BLOCK_K numKBlocks ACTIVATION idx.1 idx.2.1))))
Referenced model/specification definitions: Python source |
|---|
matmul_triton_autotunematmul_autotune_io_correctnessAbstract cast/store rounding | Configured model / stage Configured matmul/activation branches and output layout. The theorem quantifies over its named parameters but does not verify the Python autotuner or every autotune candidate. | Statement and preconditionsspecification matmul_autotune_io_correctness (R : RoundingModel)
(A B C : RegionName)
(M N K sam sak sbk sbn scm scn BM BN BK GM numKBlocks : Nat) (ACTIVATION : Bool)
(hK : K = BK * numKBlocks) (hcn : scn = 1) (hBN : BN ≤ scm) :
matmulAutotuneIO A B C M N K sam sak sbk sbn scm scn BM BN BK GM numKBlocks ACTIVATION
⊨[R] fun _ _ xs ys l =>
act ACTIVATION (∑ t : Fin numKBlocks, ∑ e : Fin BK,
xs t (aLane BM BN BK l e) * ys t (bLane BM BN BK l e))
IO definition: matmulAutotuneIOdef matmulAutotuneIO (A B C : RegionName)
(M N K sam sak sbk sbn scm scn BM BN BK GM numKBlocks : Nat) (ACTIVATION : Bool) :
StreamMasked2DKernelIO₂ where
kernel := matmul_autotune_surface A B C M N K sam sak sbk sbn scm scn BM BN BK GM
numKBlocks ACTIVATION
inp1 := A
inp2 := B
out := C
T := numKBlocks
B1 := BM * BK
B2 := BK * BN
C := BM * BN
outDType := .fp16
read1 := fun p₀ _ t l =>
((pidM p₀ M N BM BN GM * BM + l.val / BK) % M) * sam + (t.val * BK + l.val % BK) * sak
read2 := fun p₀ _ t l =>
(t.val * BK + l.val / BN) * sbk + ((pidN p₀ M N BM BN GM * BN + l.val % BN) % N) * sbn
write := fun p₀ _ l =>
scm * (pidM p₀ M N BM BN GM * BM + l.val / BN) + scn * (pidN p₀ M N BM BN GM * BN + l.val % BN)
mask1 := fun _ _ t l => t.val * BK + l.val % BK < K
mask2 := fun _ _ t l => t.val * BK + l.val / BN < K
writeMask := fun p₀ _ l =>
pidM p₀ M N BM BN GM * BM + l.val / BN < M ∧ pidN p₀ M N BM BN GM * BN + l.val % BN < N
Referenced model/specification definitions: Python source |
|---|
matrix_reductionload_reduce_kernel_output_summaryMathematical execution | Configured model / stage One declared matrix tile reduced to the displayed output addresses; positive reduction width is explicit in the IO contract. The theorem does not infer a surrounding host reduction. | Statement and preconditionsspecification load_reduce_kernel_output_summary
(x_ptr y_ptr : RegionName)
(stride_xm stride_xn stride_y BLOCK_M BLOCK_N : Nat)
(s : BlockState) :
(∃ alg, (load_reduce_kernel x_ptr y_ptr stride_xm stride_xn stride_y
BLOCK_M BLOCK_N).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := load_reduce_kernel x_ptr y_ptr stride_xm stride_xn stride_y
BLOCK_M BLOCK_N)
(initialState := s)
(write := fun i : Fin BLOCK_M => some (y_ptr, i.val))
(expected := fun i =>
matrixReduceSpec s x_ptr stride_xm stride_xn BLOCK_M BLOCK_N i)
Referenced model/specification definitions: Python source |
|---|
matrix_reductionmatrix_reduction_io_correctnessMathematical execution | Configured model / stage One declared matrix tile reduced to the displayed output addresses; positive reduction width is explicit in the IO contract. The theorem does not infer a surrounding host reduction. | Statement and preconditionsspecification matrix_reduction_io_correctness (x_ptr y_ptr : RegionName)
(stride_xm stride_xn stride_y BLOCK_M BLOCK_N : Nat) (hBN : 0 < BLOCK_N) :
matrixReduceIO x_ptr y_ptr stride_xm stride_xn stride_y BLOCK_M BLOCK_N
⊨ fun _pid xs idx => matrixReduceSpecOf BLOCK_M BLOCK_N xs idx
IO definition: matrixReduceIOdef matrixReduceIO (x_ptr y_ptr : RegionName)
(stride_xm stride_xn stride_y BLOCK_M BLOCK_N : Nat) :
MaskedTileKernelIO₁ where
kernel := load_reduce_kernel x_ptr y_ptr stride_xm stride_xn stride_y BLOCK_M
BLOCK_N
inp := x_ptr
out := y_ptr
shape := [BLOCK_M, BLOCK_N]
read := fun _pid idx => xOffset stride_xm stride_xn BLOCK_M BLOCK_N idx
write := fun _pid idx => idx.1.val
mask := fun _pid _ => True
writeMask := fun _pid idx => idx.2.1.val = 0
Referenced model/specification definitions: Python source |
|---|
matrix_transposekernel_correctnessMathematical execution | Configured model / stage One declared matrix transpose tile with injective output addresses and arbitrary valid IO placement. Hardware layout/launch selection is external. | Statement and preconditionsspecification kernel_correctness
(M Out : RegionName)
(matrix_stridex matrix_stridey out_stridex out_stridey
SIZE_M D_HEAD : Nat)
(hOutInj : Function.Injective
(fun idx : TileIndex [D_HEAD, SIZE_M] =>
outAddr out_stridex out_stridey idx)) :
transposeIO M Out matrix_stridex matrix_stridey out_stridex out_stridey
SIZE_M D_HEAD
⊨ fun _ _ xs j => xs j
IO definition: transposeIOdef transposeIO (M Out : RegionName)
(matrix_stridex matrix_stridey out_stridex out_stridey
SIZE_M D_HEAD : Nat) :
Masked2DKernelIO₁ where
kernel := kernel M Out matrix_stridex matrix_stridey out_stridex out_stridey
SIZE_M D_HEAD
inp := M
out := Out
B := D_HEAD * SIZE_M
read := fun _ _ j =>
matrixAddr matrix_stridex matrix_stridey (Lane2D.decode j)
write := fun _ _ j =>
outAddr out_stridex out_stridey (Lane2D.decode j)
mask := fun _ _ _ => True
/-- **The headline**: `kernel` implements the matrix transpose on its IO
signature — for every disjoint flat placement of `M`/`Out`, every program id
whose lanes are in bounds, and every launch state whose transposed source
window holds the tile `xs`, the translated pointer kernel terminates, output
lane `j` holds `xs j` (i.e. the output cell `(d, m)` holds the source cell
`(m, d)` — the transposition lives in the signature's windows, which this
Referenced model/specification definitions: Python source |
|---|
matrix_vector_multipmv_kernel_output_summary_generalMathematical execution | Configured model / stage One-block matrix-vector product with BLOCK_M reduction extent and explicit masks/address injectivity. This is not a proof for arbitrary multi-block reduction composition. | Statement and preconditionsspecification mv_kernel_output_summary_general
(A B C : RegionName)
(N M stride_an stride_am stride_bm stride_cn BLOCK_N BLOCK_M : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_N => cOffset s stride_cn BLOCK_N i)) :
mv_kernel_general_prop A B C N M stride_an stride_am stride_bm stride_cn
BLOCK_N BLOCK_M s
Referenced model/specification definitions: Python source |
|---|
matrix_vector_multipmv_one_block_io_correctnessMathematical execution | Configured model / stage One-block matrix-vector product with BLOCK_M reduction extent and explicit masks/address injectivity. This is not a proof for arbitrary multi-block reduction composition. | Statement and preconditionsspecification mv_one_block_io_correctness (A B C : RegionName)
(N M stride_an stride_am stride_bm stride_cn BLOCK_N BLOCK_M : Nat)
(hBN : 0 < BLOCK_N) (hBM : 0 < BLOCK_M)
(hOutInj : ∀ p₀ : Nat, Function.Injective
(fun i : Fin BLOCK_N => (p₀ * BLOCK_N + i.val) * stride_cn)) :
mvOneBlockIO A B C N M stride_an stride_am stride_bm stride_cn BLOCK_N
BLOCK_M
⊨ fun p₀ _p₁ xs ys k =>
mvSpecOf N M BLOCK_N BLOCK_M p₀ xs ys k.1
IO definition: mvOneBlockIOdef mvOneBlockIO (A B C : RegionName)
(N M stride_an stride_am stride_bm stride_cn BLOCK_N BLOCK_M : Nat) :
MaskedTile2DKernelIO₂ where
kernel := mv_kernel_one_block A B C N M stride_an stride_am stride_bm
stride_cn BLOCK_N BLOCK_M
in1 := A
in2 := B
out := C
shape := [BLOCK_N, BLOCK_M]
read1 := fun p₀ _p₁ k =>
(p₀ * BLOCK_N + k.1.val) * stride_an + k.2.1.val * stride_am
read2 := fun _p₀ _p₁ k => k.2.1.val * stride_bm
write := fun p₀ _p₁ k => (p₀ * BLOCK_N + k.1.val) * stride_cn
mask := fun p₀ _p₁ k => p₀ * BLOCK_N + k.1.val < N ∧ k.2.1.val < M
read2Mask := fun _p₀ _p₁ k => k.2.1.val < M
writeMask := fun p₀ _p₁ k => p₀ * BLOCK_N + k.1.val < N ∧ k.2.1.val = 0
Referenced model/specification definitions: Python source |
|---|
max_reductionmax_kernel_output_summaryMathematical execution | Configured model / stage The summary covers the modeled value/index dimension reduction; the IO theorem separately covers only first-stage max_kernel_1. These stages are not automatically composed. | Statement and preconditionsspecification max_kernel_output_summary
(inp out_value out_index : RegionName)
(M N K BLOCK_M BLOCK_N : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_M => maxKernelOutOffset s K BLOCK_M i))
(hOutRegions : out_value ≠ out_index) :
(∃ alg, (max_kernel inp out_value out_index M N K BLOCK_M BLOCK_N).toAlgorithm?
= Except.ok alg) ∧
ComputeCorrect.OutputPairWhere
(max_kernel inp out_value out_index M N K BLOCK_M BLOCK_N)
s out_value out_index
(maxKernelOutOffset s K BLOCK_M)
(fun i : Fin BLOCK_M => s.pids 0 * BLOCK_M + i.val < M)
(maxKernelValueSpec s inp M N K BLOCK_M BLOCK_N)
(maxKernelIndexSpec s inp M N K BLOCK_M BLOCK_N)
Referenced model/specification definitions: Python source |
|---|
max_reductionmax_kernel_1_io_correctnessMathematical execution | Configured model / stage The summary covers the modeled value/index dimension reduction; the IO theorem separately covers only first-stage max_kernel_1. These stages are not automatically composed. | Statement and preconditionsspecification max_kernel_1_io_correctness
(inp mid : RegionName) (M BLOCK_SIZE : Nat) (hB : 0 < BLOCK_SIZE) :
maxKernel1IO inp mid M BLOCK_SIZE
⊨ fun pid xs _ => maxTileSpecOf M BLOCK_SIZE pid xs
IO definition: maxKernel1IOdef maxKernel1IO (inp mid : RegionName) (M BLOCK_SIZE : Nat) :
MaskedTileKernelIO₁ where
kernel := max_kernel_1 inp mid M BLOCK_SIZE
inp := inp
out := mid
shape := [BLOCK_SIZE]
read := fun pid idx => pid * BLOCK_SIZE + idx.1.val
write := fun pid _ => pid
mask := fun pid idx => pid * BLOCK_SIZE + idx.1.val < M
writeMask := fun _ idx => idx.1.val = 0
Referenced model/specification definitions: Python source |
|---|
mean_reductionmean_dim_kernel_output_summaryMathematical execution | Configured model / stage Full modeled row mean through positive-width reduction blocks, with separate value and streaming IO contracts. Mathematical division and the declared memory layout define the scope. | Statement and preconditionsspecification mean_dim_kernel_output_summary
(X Mean : RegionName)
(M N BLOCK_M BLOCK_N : Nat) (s : BlockState)
(hStepNe : BLOCK_N ≠ 0) :
(∃ alg, (mean_dim_kernel X Mean M N BLOCK_M BLOCK_N).toAlgorithm? =
Except.ok alg) ∧
mean_dim_kernel_correct_target X Mean M N BLOCK_M BLOCK_N s
Referenced model/specification definitions: Python source |
|---|
mean_reductionmean_dim_kernel_io_correctnessAbstract cast/store rounding | Configured model / stage Full modeled row mean through positive-width reduction blocks, with separate value and streaming IO contracts. Mathematical division and the declared memory layout define the scope. | Statement and preconditionsspecification mean_dim_kernel_io_correctness (R : RoundingModel)
(X Mean : RegionName) (M N BLOCK_M BLOCK_N : Nat)
(hStepNe : BLOCK_N ≠ 0) :
meanKernelIO X Mean M N BLOCK_M BLOCK_N ⊨[R] fun _ _ xs i =>
meanStreamSpec N BLOCK_M BLOCK_N xs i
IO definition: meanKernelIOdef meanKernelIO (X Mean : RegionName) (M N BLOCK_M BLOCK_N : Nat) :
StreamMasked2DKernelIO₁ where
kernel := mean_dim_kernel X Mean M N BLOCK_M BLOCK_N
inp1 := X
out := Mean
T := meanNumSteps N BLOCK_N
B1 := BLOCK_M * BLOCK_N
C := BLOCK_M
read1 := fun p₀ _ t l => meanReadAddr N BLOCK_M BLOCK_N p₀ t.val l
write := fun p₀ _ i => meanWriteAddr BLOCK_M p₀ i
mask1 := fun p₀ _ t l => meanReadActive M N BLOCK_M BLOCK_N p₀ t.val l
writeMask := fun p₀ _ i => meanWriteActive M BLOCK_M p₀ i
Referenced model/specification definitions: Python source |
|---|
mixed_sparse_attentionmixed_sparse_attention_output_closed_form_summary_generalMathematical execution | Configured model / stage Forward sparse-layout value facts under supplied block/column metadata and the stated bounds; read the referenced sparse specification for the selected support. | Statement and preconditionsspecification mixed_sparse_attention_output_closed_form_summary_general
(Q K V Out : RegionName)
(Seqlens Blocks BlockOffsets ColCounts Cols : Region .nat) (s : BlockState)
(BM BN BD : Nat) (hBN : 0 < BN) (hBN16 : 16 ≤ BN)
-- memory-layout strides (Q/K/V/O batch z, head h, row m / key n, channel k)
-- and grid/layout sizes, ALL symbolic
(NCTX Zc H NR NS NV
sqz sqh sqm sqk skz skh skn skk svz svh svn svk soz soh som sok : Nat)
-- honest contiguity hypotheses (the natural row-major attention layout):
-- channel strides are 1; V reuses K's batch/head base; O channel stride 1 and
-- its row stride covers the channel block (BD ≤ stride_om)
(hsqk : sqk = 1) (hskk : skk = 1) (hsvk : svk = 1) (hvz : svz = skz) (hvh : svh = skh)
(hsok : sok = 1) (hBDsom : BD ≤ som)
(sm_scale : ℝ)
(hundef : ∀ rg o, s.undef rg o = 0)
(hactive : s.pids 0 * BM < seqLen s H (Region.cast Seqlens))
(hNCBN : s.readMemValue .nat (Region.cast ColCounts) (s.pids 1 * NR + s.pids 0) ≤ BN)
(hpos : ∀ i : Fin BM, s.pids 0 * BM + i.val < seqLen s H (Region.cast Seqlens) →
0 < msaDenomUpto BM BN
(msaCatScore0GS Q K V Seqlens Blocks BlockOffsets ColCounts Cols BM BN BD
H NR NS NV sqz sqh sqm sqk skz skh skn skk s sm_scale) 9 i) :
∃ sF, exec (mixed_sparse_attention_fwd_kernel_surface Q K V Seqlens
sm_scale Blocks BlockOffsets ColCounts Cols Out
sqz sqh sqm sqk skz skh skn skk svz svh svn svk soz soh som sok
Zc H NCTX NR NS NV BM BN BD FloatDType.fp16).toAlgKernel s = some sF
∧ ∀ idx : TileIndex [BM, BD],
active s H Seqlens BM idx →
sF.readMemValue .fp16 Out (outOffset s H sqz sqh som sok BM idx)
= (some (mixedSparseAttnClosedForm s Q K V BlockOffsets Cols H
sqz sqh sqm skz skh skn svz svh svn NR NS NV
(s.readMemValue .nat (Region.cast Blocks) (s.pids 1 * NR + s.pids 0))
(s.readMemValue .nat (Region.cast ColCounts) (s.pids 1 * NR + s.pids 0))
(seqLen s H (Region.cast Seqlens)) BD BM BN sm_scale idx.1 (dIndex idx)) : WithBot ℝ)
Referenced model/specification definitions: Python source |
|---|
mixed_sparse_attentionmixed_sparse_attention_epilogue_io_correctnessMathematical execution | Precomputed-input slice Epilogue divides a supplied AccPre tile by supplied LPre normalizers and stores active outputs; sparse attention production is outside this IO kernel. | Statement and preconditionsspecification mixed_sparse_attention_epilogue_io_correctness
(AccPre LPre : RegionName) (Seqlens : Region .nat) (Out : RegionName)
(H stride_qz stride_qh stride_om stride_ok BLOCK_M BLOCK_DMODEL : Nat)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
p₁ / H * stride_qz + p₁ % H * stride_qh +
(p₀ * BLOCK_M + idx.1.val) * stride_om + idx.2.1.val * stride_ok)) :
epilogueIO AccPre LPre Seqlens Out H stride_qz stride_qh stride_om stride_ok
BLOCK_M BLOCK_DMODEL
⊨ fun _p₀ _p₁ _m xs ys idx => xs idx / ys (idx.1, PUnit.unit)
IO definition: epilogueIOdef epilogueIO (AccPre LPre : RegionName) (Seqlens : Region .nat)
(Out : RegionName)
(H stride_qz stride_qh stride_om stride_ok BLOCK_M BLOCK_DMODEL : Nat) :
Meta1MaskedTileShapedKernelIO₂ where
kernel := mixed_sparse_attention_epilogue_slice AccPre LPre Seqlens Out H
stride_qz stride_qh stride_om stride_ok BLOCK_M BLOCK_DMODEL
mbuf := Region.cast Seqlens
in1 := AccPre
in2 := LPre
out := Out
shape1 := [BLOCK_M, BLOCK_DMODEL]
shape2 := [BLOCK_M]
shapeOut := [BLOCK_M, BLOCK_DMODEL]
mwin := fun _p₀ p₁ => p₁ / H
read1 := fun p₀ p₁ _m idx =>
p₁ / H * stride_qz + p₁ % H * stride_qh +
(p₀ * BLOCK_M + idx.1.val) * stride_om + idx.2.1.val * stride_ok
read2 := fun p₀ _p₁ _m i => p₀ * BLOCK_M + i.1.val
write := fun p₀ p₁ _m idx =>
p₁ / H * stride_qz + p₁ % H * stride_qh +
(p₀ * BLOCK_M + idx.1.val) * stride_om + idx.2.1.val * stride_ok
mask1 := fun _p₀ _p₁ _m _idx => True
mask2 := fun _p₀ _p₁ _m _i => True
writeMask := fun p₀ _p₁ m idx => p₀ * BLOCK_M + idx.1.val < m
Referenced model/specification definitions: Python source |
|---|
mul_exponent_compensatormul_kernel_correctnessMathematical execution | Configured model / stage One modeled positive-size tile multiplied by the declared exponent compensator. This does not certify hardware exponent-bit manipulation or a larger scaling pipeline. | Statement and preconditionsspecification mul_kernel_correctness (src dst : RegionName) (B : Nat)
(hB : 0 < B) :
mulIO src dst B ⊨ fun xs i => xs i * exponentCompensator
IO definition: mulIOnoncomputable def mulIO (src dst : RegionName) (B : Nat) : KernelIO₁ where
kernel := mul_kernel src dst B
inp := src
out := dst
Bin := B
Bout := B
read := fun pid => pid * B
write := fun pid => pid * B
Referenced model/specification definitions: Python source |
|---|
nested_loops_processingnested3_output_summary_generalMathematical execution | Configured model / stage Facts for explicitly selected nested-loop stores and shifted 2x2 copies. The IO headlines cover the first a1 store only, not the whole nested-loop program. | Statement and preconditionsspecification nested3_output_summary_general
(in_ptr out_ptr : RegionName) (stride_m stride_n : Nat)
(s : BlockState)
(hInj0 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffset stride_m stride_n idx))
(hInj2 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 2 idx))
(hInj4 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 4 idx))
(hInj6 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 6 idx))
(hInj8 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 8 idx))
(hInj10 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 10 idx))
(hInj12 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 12 idx))
(hInj14 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 14 idx))
(hInj16 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 16 idx))
(hInj18 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 18 idx))
(hInj20 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 20 idx))
(hInj22 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 22 idx))
(hInj24 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 24 idx))
(hInj26 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 26 idx))
(hInj28 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 28 idx))
(hInj30 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 30 idx))
(hInj32 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 32 idx))
(hInj34 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 34 idx))
(hInj36 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 36 idx))
(hInj38 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 38 idx))
(hInj40 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 40 idx))
(hInj42 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 42 idx))
(hInj44 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 44 idx))
(hInj46 : Function.Injective
(fun idx : TileIndex [2, 2] => matrixOffsetShift stride_m stride_n 46 idx)) :
-- (a) The full nested-loop surface lowers to the algorithm layer.
(∃ alg, (nested3 in_ptr out_ptr stride_m stride_n).toAlgorithm? =
Except.ok alg) ∧
-- (b1) Initial a1 store: verbatim copy at shift 0.
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_first_a1_store in_ptr out_ptr stride_m stride_n)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffset stride_m stride_n idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffset stride_m stride_n idx)) ∧
-- (b2) First a2 store (i,j,k)=(0,0,0).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_first_a2_store in_ptr out_ptr stride_m stride_n)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 2 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 2 idx)) ∧
-- (b3) First a3 store (i,j,k)=(0,0,0).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_first_a3_store in_ptr out_ptr stride_m stride_n)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 4 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 4 idx)) ∧
-- (b4) Second k: a1 reused (in 0 → out 6).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 0 6)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 6 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 0 idx)) ∧
-- (b5) Second k: a2 (in 2 → out 8).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 2 8)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 8 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 2 idx)) ∧
-- (b6) Second k: a3 (in 6 → out 10).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 6 10)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 10 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 6 idx)) ∧
-- (b7) Second j, first k: a1 reused (in 0 → out 12).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 0 12)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 12 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 0 idx)) ∧
-- (b8) Second j, first k: a2 (in 8 → out 14).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 8 14)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 14 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 8 idx)) ∧
-- (b9) Second j, first k: a3 (in 10 → out 16).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 10 16)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 16 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 10 idx)) ∧
-- (b10) Second j, second k: a1 reused (in 0 → out 18).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 0 18)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 18 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 0 idx)) ∧
-- (b11) Second j, second k: a2 reused (in 8 → out 20).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 8 20)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 20 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 8 idx)) ∧
-- (b12) Second j, second k: a3 (in 12 → out 22).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 12 22)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 22 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 12 idx)) ∧
-- (b13) Second i, first j, first k: a1 (in 14 → out 24).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 14 24)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 24 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 14 idx)) ∧
-- (b14) Second i, first j, first k: a2 (in 16 → out 26).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 16 26)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 26 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 16 idx)) ∧
-- (b15) Second i, first j, first k: a3 (in 18 → out 28).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 18 28)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 28 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 18 idx)) ∧
-- (b16) Second i, first j, second k: a1 reused (in 14 → out 30).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 14 30)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 30 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 14 idx)) ∧
-- (b17) Second i, first j, second k: a2 reused (in 16 → out 32).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 16 32)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 32 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 16 idx)) ∧
-- (b18) Second i, first j, second k: a3 (in 20 → out 34).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 20 34)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 34 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 20 idx)) ∧
-- (b19) Second i, second j, first k: a1 reused (in 14 → out 36).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 14 36)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 36 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 14 idx)) ∧
-- (b20) Second i, second j, first k: a2 (in 22 → out 38).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 22 38)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 38 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 22 idx)) ∧
-- (b21) Second i, second j, first k: a3 (in 24 → out 40).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 24 40)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 40 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 24 idx)) ∧
-- (b22) Second i, second j, second k: a1 reused (in 14 → out 42).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 14 42)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 42 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 14 idx)) ∧
-- (b23) Second i, second j, second k: a2 reused (in 22 → out 44).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 22 44)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 44 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 22 idx)) ∧
-- (b24) Second i, second j, second k: final a3 (in 26 → out 46).
ComputeCorrect.Realizes_without_Rounding
(kernel := nested3_shifted_copy_store in_ptr out_ptr stride_m stride_n 26 46)
(initialState := s)
(write := fun idx : TileIndex [2, 2] =>
some (out_ptr, matrixOffsetShift stride_m stride_n 46 idx))
(expected := fun idx =>
s.readMem in_ptr (matrixOffsetShift stride_m stride_n 26 idx))
Referenced model/specification definitions: Python source |
|---|
nested_loops_processingnested3_first_a1_store_io_correctnessMathematical execution | Configured model / stage Facts for explicitly selected nested-loop stores and shifted 2x2 copies. The IO headlines cover the first a1 store only, not the whole nested-loop program. | Statement and preconditionsspecification nested3_first_a1_store_io_correctness
(in_ptr out_ptr : RegionName) (stride_m stride_n : Nat)
(hOutInj : Function.Injective
(fun idx : TileIndex [2, 2] =>
idx.1.val * stride_m + idx.2.1.val * stride_n)) :
a1StoreIO in_ptr out_ptr stride_m stride_n ⊨ fun _pid xs idx => xs idx
IO definition: a1StoreIOdef a1StoreIO (in_ptr out_ptr : RegionName) (stride_m stride_n : Nat) :
MaskedTileKernelIO₁ where
kernel := nested3_first_a1_store in_ptr out_ptr stride_m stride_n
inp := in_ptr
out := out_ptr
shape := [2, 2]
read := fun _pid idx => idx.1.val * stride_m + idx.2.1.val * stride_n
write := fun _pid idx => idx.1.val * stride_m + idx.2.1.val * stride_n
mask := fun _pid _idx => True
Referenced model/specification definitions: Python source |
|---|
nested_loops_processingnested3_first_a1_store_io_correctnessRAbstract cast/store rounding | Configured model / stage Facts for explicitly selected nested-loop stores and shifted 2x2 copies. The IO headlines cover the first a1 store only, not the whole nested-loop program. | Statement and preconditionsspecification nested3_first_a1_store_io_correctnessR (R : RoundingModel)
(in_ptr out_ptr : RegionName) (stride_m stride_n : Nat)
(hOutInj : Function.Injective
(fun idx : TileIndex [2, 2] =>
idx.1.val * stride_m + idx.2.1.val * stride_n)) :
a1StoreIO in_ptr out_ptr stride_m stride_n
⊨[R, FloatDType.real] fun _pid xs idx => xs idx
IO definition: a1StoreIOdef a1StoreIO (in_ptr out_ptr : RegionName) (stride_m stride_n : Nat) :
MaskedTileKernelIO₁ where
kernel := nested3_first_a1_store in_ptr out_ptr stride_m stride_n
inp := in_ptr
out := out_ptr
shape := [2, 2]
read := fun _pid idx => idx.1.val * stride_m + idx.2.1.val * stride_n
write := fun _pid idx => idx.1.val * stride_m + idx.2.1.val * stride_n
mask := fun _pid _idx => True
Referenced model/specification definitions: Python source |
|---|
parallel_attentionpa_fwd_o_exec_genuineMathematical execution | Configured model / stage Forward O/Z and backward DK/DV stages at the displayed shape/layout configurations. The DQ stage and whole host pipeline are external; exec-existentials omit a frame. | Statement and preconditionsspecification pa_fwd_o_exec_genuine
(s : BlockState) (q k v o z : RegionName)
(s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d : Nat) (scale : ℝ)
(B H T K V BTL BTS BK BV : Nat)
(hOZ : o ≠ z) (hSd : s_vo_d = 1) (hσ : BV ≤ s_vo_t)
(hBTS : BTL % BTS = 0) (hBTSpos : 0 < BTS) :
∃ sF, exec (pa_fwd_surface q k v o z s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t
s_vo_d scale B H T K V BTL BTS BK BV).toAlgKernel s = some sF
∧ (∀ idx : TileIndex [BTL, BV], paOActive s T V BV BTL idx →
sF.readMem o (paOOffset s B H s_vo_h s_vo_t s_vo_d V BV BTL idx)
= paOOut s q k v s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d scale
T K V BTL BK BV idx.1.val idx.2.1.val)
∧ (∀ a : Fin BTL, s.pids 1 * BTL + a.val < T →
sF.readMem z (paZOffset s B H T V BV BTL a)
= paZOut s q k s_qk_h s_qk_t s_qk_d scale T K V BTL BK BV a.val)
Referenced model/specification definitions: Python source |
|---|
parallel_attentionpa_bwd_dkv_exec_genuineMathematical execution | Configured model / stage Forward O/Z and backward DK/DV stages at the displayed shape/layout configurations. The DQ stage and whole host pipeline are external; exec-existentials omit a frame. | Statement and preconditionsspecification pa_bwd_dkv_exec_genuine
(s : BlockState) (q k v do_ dz dk dv : RegionName)
(i_bh i_c i_k i_v i_h : Nat)
(s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d : Nat) (scale : ℝ)
(B H T K V BTL BTS BK BV : Nat)
(hDkDv : dk ≠ dv)
(hSkd : s_k_d = 1) (hSvd : s_v_d = 1)
(hσk : BK ≤ s_k_t) (hσv : BV ≤ s_v_t)
(hBTS : BTL % BTS = 0) (hBTSpos : 0 < BTS)
(hundef : ∀ rg off, s.undef rg off = 0) :
∃ sF, exec (pa_bwd_dkv_surface q k v do_ dz dk dv i_bh i_c i_k i_v i_h
s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d scale
B H T K V BTL BTS BK BV).toAlgKernel s = some sF
∧ (∀ idx : TileIndex [BTL, BK], pbDkActive i_c i_k T K BTL BK idx →
sF.readMem dk
(pbDkOffset i_bh i_c i_k i_v B H s_k_h s_k_t s_k_d BTL BK idx)
= pbDkOut s q k v do_ dz s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d i_bh
i_c i_k i_v scale T K V BTL BTS BK BV idx.1.val idx.2.1.val)
∧ (∀ idx : TileIndex [BTL, BV], pbDvActive i_c i_v T V BTL BV idx →
sF.readMem dv
(pbDvOffset i_bh i_c i_k i_v B H s_v_h s_v_t s_v_d BTL BV idx)
= pbDvOut s q k do_ s_k_h s_k_t s_k_d s_v_h s_v_t s_v_d i_bh i_c
i_k i_v scale T K V BTL BTS BK BV idx.1.val idx.2.1.val)
Referenced model/specification definitions: Python source |
|---|
parallel_retention_attentionpra_fwd_o_exec_genuineMathematical execution | Configured model / stage Forward O and backward DK/DV stages for explicit causal/retention and layout parameters. No DQ-stage composition or frame conjunct is established by these headlines. | Statement and preconditionsspecification pra_fwd_o_exec_genuine
(s : BlockState) (q k v o : RegionName)
(s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d : Nat) (scale : ℝ)
(B H T K V BTL BTS BK BV : Nat)
(hSd : s_vo_d = 1) (hσ : BV ≤ s_vo_t)
(hBTS : BTL % BTS = 0) (hBTSpos : 0 < BTS) :
∃ sF, exec (pra_fwd_surface q k v o s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t
s_vo_d scale B H T K V BTL BTS BK BV).toAlgKernel s = some sF
∧ (∀ idx : TileIndex [BTL, BV], praOActive s T V BV BTL idx →
sF.readMem o (praOOffset s B H s_vo_h s_vo_t s_vo_d V BV BTL idx)
= praOOut s q k v s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d scale
H T K V BTL BK BV idx.1.val idx.2.1.val)
Referenced model/specification definitions: Python source |
|---|
parallel_retention_attentionpra_bwd_dkv_exec_genuineMathematical execution | Configured model / stage Forward O and backward DK/DV stages for explicit causal/retention and layout parameters. No DQ-stage composition or frame conjunct is established by these headlines. | Statement and preconditionsspecification pra_bwd_dkv_exec_genuine
(s : BlockState) (q k v do_ dk dv : RegionName)
(i_bh i_c i_k i_v i_h : Nat)
(s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d : Nat) (scale : ℝ)
(B H T K V BTL BTS BK BV : Nat)
(hDkDv : dk ≠ dv) (hSkd : s_qk_d = 1) (hSvd : s_vo_d = 1)
(hσk : BK ≤ s_qk_t) (hσv : BV ≤ s_vo_t)
(hBTS : BTL % BTS = 0) (hBTSpos : 0 < BTS) :
∃ sF, exec (pra_bwd_dkv_surface q k v do_ dk dv i_bh i_c i_k i_v i_h
s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d scale
B H T K V BTL BTS BK BV).toAlgKernel s = some sF
∧ (∀ idx : TileIndex [BTL, BK], prbDkActive i_c i_k T K BTL BK idx →
sF.readMem dk
(prbDkOffset i_bh i_c i_k i_v B H s_qk_h s_qk_t s_qk_d BTL BK idx)
= prbDkOut s q k v do_ s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d
i_bh i_c i_k i_v i_h scale T K V BTL BTS BK BV
idx.1.val idx.2.1.val)
∧ (∀ idx : TileIndex [BTL, BV], prbDvActive i_c i_v T V BTL BV idx →
sF.readMem dv
(prbDvOffset i_bh i_c i_k i_v B H s_vo_h s_vo_t s_vo_d BTL BV idx)
= prbDvOut s q k v do_ s_qk_h s_qk_t s_qk_d s_vo_h s_vo_t s_vo_d
i_bh i_c i_k i_v i_h scale T K V BTL BTS BK BV
idx.1.val idx.2.1.val)
Referenced model/specification definitions: Python source |
|---|
pow_scalar_tensorpow_scalar_tensor_output_summary_generalMathematical execution | Configured model / stage Rank-one scalar-base tensor-exponent model with one-tile and grid-stride value facts. The IO headline covers one tile; Real.rpow is mathematical, not a hardware approximation. | Statement and preconditionsspecification pow_scalar_tensor_output_summary_general
(val0 : ℝ) (in0_ptr out0_ptr : RegionName)
(in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0 : Nat)
(s : BlockState)
(hStride : 0 < out0_stride0)
(hDisj : in0_ptr ≠ out0_ptr)
(hGrid : 0 < s.numPids 0) :
-- (1) both branch surfaces lower to the algorithm layer
(∃ alg, (pow_func_scalar_tensor_kernel_rank_1_one_tile_surface val0
in0_ptr out0_ptr in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta
tile_size0).toAlgorithm? = Except.ok alg) ∧
(∃ alg, (pow_func_scalar_tensor_kernel_rank_1_grid_stride_surface val0
in0_ptr out0_ptr in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta
tile_size0).toAlgorithm? = Except.ok alg) ∧
-- (2) one_tile_per_cta = true: genuine elementwise scalar-base power
ComputeCorrect.Realizes_without_Rounding
(kernel := pow_func_scalar_tensor_kernel_rank_1_one_tile_surface val0
in0_ptr out0_ptr in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta
tile_size0)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin tile_size0 => taskIndex (s.pids 0) tile_size0 i < s0)
(fun i => (out0_ptr, taskIndex (s.pids 0) tile_size0 i * out0_stride0)))
(expected := fun i =>
powSpec s in0_ptr in0_stride0 val0
(taskIndex (s.pids 0) tile_size0 i)) ∧
-- (3) one_tile_per_cta = false: genuine elementwise scalar-base power
-- across the whole grid-stride loop
ComputeCorrect.Realizes_without_Rounding
(kernel := pow_func_scalar_tensor_kernel_rank_1_grid_stride_surface val0
in0_ptr out0_ptr in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta
tile_size0)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun p : Fin tiles_per_cta × Fin tile_size0 =>
taskIndex (s.pids 0 + p.1.val * s.numPids 0) tile_size0 p.2 < s0)
(fun p => (out0_ptr,
taskIndex (s.pids 0 + p.1.val * s.numPids 0) tile_size0 p.2
* out0_stride0)))
(expected := fun p =>
powSpec s in0_ptr in0_stride0 val0
(taskIndex (s.pids 0 + p.1.val * s.numPids 0) tile_size0 p.2))
Referenced model/specification definitions: Python source |
|---|
pow_scalar_tensorpow_scalar_tensor_one_tile_io_correctnessMathematical execution | Configured model / stage Rank-one scalar-base tensor-exponent model with one-tile and grid-stride value facts. The IO headline covers one tile; Real.rpow is mathematical, not a hardware approximation. | Statement and preconditionsspecification pow_scalar_tensor_one_tile_io_correctness (val0 : ℝ)
(in0_ptr out0_ptr : RegionName)
(in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0 : Nat)
(hStride : 0 < out0_stride0) :
powOneTileIO val0 in0_ptr out0_ptr in0_stride0 out0_stride0 s0 num_tasks
tiles_per_cta tile_size0
⊨ fun _pid xs i => Real.rpow val0 (xs i)
IO definition: powOneTileIOdef powOneTileIO (val0 : ℝ) (in0_ptr out0_ptr : RegionName)
(in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0 : Nat) :
MaskedTileKernelIO₁ where
kernel := pow_func_scalar_tensor_kernel_rank_1_one_tile_surface val0 in0_ptr
out0_ptr in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0
inp := in0_ptr
out := out0_ptr
shape := [tile_size0]
read := fun pid i => taskIndex pid tile_size0 i.1 * in0_stride0
write := fun pid i => taskIndex pid tile_size0 i.1 * out0_stride0
mask := fun pid i => taskIndex pid tile_size0 i.1 < s0
Referenced model/specification definitions: Python source |
|---|
quant_transpose_kernelquantize_global_transpose_blocked_output_summary_generalMathematical execution | Blocked surface The faithful quantized transpose has a declared fixed-width rounding/projection blocker; the companion value result describes a scaled real store slice. | Statement and preconditionsspecification quantize_global_transpose_blocked_output_summary_general
(A AbsmaxInv B : RegionName)
(stride_am stride_an stride_bn stride_bm M N BLOCK_M BLOCK_N GROUP_M : Nat)
(scale127 : ℝ)
(s : BlockState)
(hOutInj : Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
bOffset s stride_bm stride_bn BLOCK_M BLOCK_N idx)) :
(∃ err, (quantize_global_transpose_real_surface A AbsmaxInv B
stride_am stride_an stride_bn stride_bm M N BLOCK_M BLOCK_N
GROUP_M).toAlgorithm? = Except.error err) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := quantize_global_transpose_scaled_store_slice A AbsmaxInv B
stride_am stride_an stride_bn stride_bm M N BLOCK_M BLOCK_N scale127)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BLOCK_M BLOCK_N)
(fun idx => (B, bOffset s stride_bm stride_bn BLOCK_M BLOCK_N idx)))
(expected := fun idx =>
quantTransposeScaledSpec s A AbsmaxInv stride_am stride_an
BLOCK_M BLOCK_N scale127 idx)
Referenced model/specification definitions: Python source |
|---|
quant_transpose_kernelquant_transpose_scaled_store_io_correctnessMathematical execution | Pre-rounding slice Transposed scaled real-valued store using supplied AbsmaxInv; llrint/int8 rounding and any scale producer are outside the contract. | Statement and preconditionsspecification quant_transpose_scaled_store_io_correctness
(A AbsmaxInv B : RegionName)
(stride_am stride_an stride_bn stride_bm M N BLOCK_M BLOCK_N : Nat)
(scale127 : ℝ)
(hM : 0 < BLOCK_M) (hN : 0 < BLOCK_N)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
(p₀ * BLOCK_M + idx.1.val) * stride_bm
+ (p₁ * BLOCK_N + idx.2.1.val) * stride_bn)) :
quantTransposeScaledIO A AbsmaxInv B stride_am stride_an stride_bn stride_bm
M N BLOCK_M BLOCK_N scale127
⊨ fun _p₀ _p₁ xs ys idx => scale127 * (xs idx * ys idx)
IO definition: quantTransposeScaledIOdef quantTransposeScaledIO (A AbsmaxInv B : RegionName)
(stride_am stride_an stride_bn stride_bm M N BLOCK_M BLOCK_N : Nat)
(scale127 : ℝ) : MaskedTile2DKernelIO₂ where
kernel := quantize_global_transpose_scaled_store_slice A AbsmaxInv B
stride_am stride_an stride_bn stride_bm M N BLOCK_M BLOCK_N scale127
in1 := A
in2 := AbsmaxInv
out := B
shape := [BLOCK_M, BLOCK_N]
read1 := fun p₀ p₁ idx =>
(p₀ * BLOCK_M + idx.1.val) * stride_am
+ (p₁ * BLOCK_N + idx.2.1.val) * stride_an
read2 := fun _ _ _ => 0
write := fun p₀ p₁ idx =>
(p₀ * BLOCK_M + idx.1.val) * stride_bm
+ (p₁ * BLOCK_N + idx.2.1.val) * stride_bn
mask := fun p₀ p₁ idx =>
p₀ * BLOCK_M + idx.1.val < M ∧ p₁ * BLOCK_N + idx.2.1.val < N
read2Mask := fun _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
quantize_copy_kvdestindex_copy_quantize_kv_output_summary_generalMathematical execution | Configured model / stage Modeled KV quantization and scale outputs with destination metadata and an abstract rounding contract. Integer/floating conversion definitions must not be read as concrete int8 overflow or encoding guarantees. | Statement and preconditionsspecification destindex_copy_quantize_kv_output_summary_general
(K DestLoc Out OutScale : RegionName)
(stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num BLOCK_DMODEL BLOCK_HEAD : Nat)
(s : BlockState) (hD : 0 < BLOCK_DMODEL) (hOut : Out ≠ OutScale)
(hValInj : Function.Injective
(fun idx : TileIndex [BLOCK_HEAD, BLOCK_DMODEL] =>
outOffset s DestLoc stride_o_bs stride_o_h stride_o_d idx))
(hScaleInj : Function.Injective
(fun i : Fin BLOCK_HEAD => scaleOutOffset1 s DestLoc stride_os_bs stride_os_h i)) :
(∃ alg, (destindex_copy_quantize_kv_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num BLOCK_DMODEL BLOCK_HEAD).toAlgorithm? = Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := destindex_copy_quantize_kv_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num BLOCK_DMODEL BLOCK_HEAD)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_HEAD, BLOCK_DMODEL] =>
some (Out, outOffset s DestLoc stride_o_bs stride_o_h stride_o_d idx))
(expected := fun idx : TileIndex [BLOCK_HEAD, BLOCK_DMODEL] =>
(if active s head_num BLOCK_HEAD BLOCK_DMODEL idx then
quantizeCopyKvSurfaceIntValue s K stride_k_bs stride_k_h stride_k_d head_num BLOCK_DMODEL hD idx
else s.readMemValue .int Out (outOffset s DestLoc stride_o_bs stride_o_h stride_o_d idx) : Int))) ∧
-- Conjunct (3): raw `.map … = some (if …)` form, NOT `Realizes_without_Rounding` — see the
-- honest carrier note in the docstring. `TileCarrier .fp16` (the decoded fp16
-- value read back here) has no `OutputReadable` instance in the framework, so
-- this genuine, non-self-referential output cannot be wrapped in `Realizes_without_Rounding`.
(∀ i : Fin BLOCK_HEAD,
(exec (destindex_copy_quantize_kv_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num BLOCK_DMODEL BLOCK_HEAD) s).map
(·.readMemValue .fp16 OutScale (scaleOutOffset1 s DestLoc stride_os_bs stride_os_h i))
= some (if scaleActive head_num BLOCK_HEAD i then
quantizeCopyKvScaleCell s K stride_k_bs stride_k_h stride_k_d head_num BLOCK_DMODEL hD i.val
else s.readMemValue .fp16 OutScale (scaleOutOffset1 s DestLoc stride_os_bs stride_os_h i)))
Referenced model/specification definitions: Python source |
|---|
quantize_copy_kvquantize_copy_kv_io_correctnessAbstract cast/store rounding | Configured model / stage Modeled KV quantization and scale outputs with destination metadata and an abstract rounding contract. Integer/floating conversion definitions must not be read as concrete int8 overflow or encoding guarantees. | Statement and preconditionsspecification quantize_copy_kv_io_correctness (R : RoundingModel)
(K : RegionName) (DestLoc : RegionName) (Out OutScale : RegionName)
(stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d
head_num BLOCK_DMODEL BLOCK_HEAD : Nat)
(hD : 0 < BLOCK_DMODEL) (hOut : Out ≠ OutScale)
(hValInj : Function.Injective
(fun idx : TileIndex [BLOCK_HEAD, BLOCK_DMODEL] =>
stride_o_h * idx.1.val + stride_o_d * idx.2.1.val))
(hScaleInj : Function.Injective
(fun i : Fin BLOCK_HEAD => stride_os_h * i.val)) :
quantizeCopyKvIO K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num BLOCK_DMODEL BLOCK_HEAD
⊨[R] fun _ _ _ xs =>
(fun j => valueXsR R BLOCK_HEAD BLOCK_DMODEL head_num hD xs j,
fun i => scaleRealXs BLOCK_HEAD BLOCK_DMODEL head_num hD xs i)
IO definition: quantizeCopyKvIOdef quantizeCopyKvIO
(K : RegionName) (DestLoc : RegionName) (Out OutScale : RegionName)
(stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d
head_num BLOCK_DMODEL BLOCK_HEAD : Nat) :
MetaMasked2DKernelIO₁ₓ₂ where
kernel := destindex_copy_quantize_kv_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num BLOCK_DMODEL BLOCK_HEAD
mbuf1 := DestLoc
inp := K
out1 := Out
out2 := OutScale
B := BLOCK_HEAD * BLOCK_DMODEL
C := BLOCK_HEAD
oty1 := .int
oty2 := .float
out2DType := .fp16
mwin1 := fun pid₀ _ => pid₀
read := fun pid₀ _ _ j =>
pid₀ * stride_k_bs + (j.val / BLOCK_DMODEL) * stride_k_h +
stride_k_d * (j.val % BLOCK_DMODEL)
write1 := fun _ _ m1 j =>
m1 * stride_o_bs + stride_o_h * (j.val / BLOCK_DMODEL) +
stride_o_d * (j.val % BLOCK_DMODEL)
write2 := fun _ _ m1 i => m1 * stride_os_bs + stride_os_h * i.val
mask := fun _ _ _ j => j.val / BLOCK_DMODEL < head_num
writeMask2 := fun _ _ _ i => i.val < head_num
Referenced model/specification definitions: Python source |
|---|
quantize_globalquantize_global_correctnessMathematical execution | Pre-rounding slice Scaled real store using a supplied inverse maximum. The headline explicitly rejects faithful llrint/int8 projection; the rounded quantized result is not proved. | Statement and preconditionsspecification quantize_global_correctness
(x_ptr absmax_inv_ptr output_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) (scale127 : ℝ)
(hBS : 0 < BLOCK_SIZE) :
(∃ err, (quantize_global_surface x_ptr absmax_inv_ptr output_ptr
n_elements BLOCK_SIZE).toAlgorithm? = Except.error err) ∧
(quantizeGlobalIO x_ptr absmax_inv_ptr output_ptr n_elements BLOCK_SIZE
scale127 ⊨ fun _ _ xs ys i => scale127 * (xs i * ys i))
IO definition: quantizeGlobalIOdef quantizeGlobalIO (x_ptr absmax_inv_ptr output_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) (scale127 : ℝ) : Masked2DKernelIO₂ where
kernel := quantize_global_scaled_store_slice x_ptr absmax_inv_ptr output_ptr
n_elements BLOCK_SIZE scale127
in1 := x_ptr
in2 := absmax_inv_ptr
out := output_ptr
B := BLOCK_SIZE
read1 := fun pid _ j => pid * BLOCK_SIZE + j.val
read2 := fun _ _ _ => 0
write := fun pid _ j => pid * BLOCK_SIZE + j.val
mask := fun pid _ j => pid * BLOCK_SIZE + j.val < n_elements
read2Mask := fun _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
quantize_kv_copydestindex_copy_quantize_kv_group_output_summary_generalMathematical execution | Configured model / stage Groupwise KV quantization and scale outputs in the mathematical conversion model. Source has repeated Python kernel names; source links identify definitions by line. No concrete int8 hardware guarantee. | Statement and preconditionsspecification destindex_copy_quantize_kv_group_output_summary_general
(K DestLoc Out OutScale : RegionName)
(stride_k_bs stride_k_h stride_k_g stride_k_d
stride_o_bs stride_o_h stride_o_g stride_o_d
stride_os_bs stride_os_h stride_os_g
group_size BLOCK_GROUP_NUM BLOCK_GROUP_DIM : Nat)
(s : BlockState) (hD : 0 < BLOCK_GROUP_DIM) (hOut : Out ≠ OutScale)
(hOutInj : Function.Injective
(fun idx : TileIndex [BLOCK_GROUP_NUM, BLOCK_GROUP_DIM] =>
outOffset s DestLoc stride_o_bs stride_o_h stride_o_g stride_o_d idx))
(hScaleInj : Function.Injective
(fun i : Fin BLOCK_GROUP_NUM =>
scaleOutOffset s DestLoc stride_os_bs stride_os_h i)) :
(∃ alg,
(destindex_copy_quantize_kv_group_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_g stride_k_d stride_o_bs stride_o_h
stride_o_g stride_o_d stride_os_bs stride_os_h stride_os_g group_size
BLOCK_GROUP_NUM BLOCK_GROUP_DIM).toAlgorithm? =
Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := destindex_copy_quantize_kv_group_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_g stride_k_d stride_o_bs stride_o_h
stride_o_g stride_o_d stride_os_bs stride_os_h stride_os_g group_size
BLOCK_GROUP_NUM BLOCK_GROUP_DIM)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_GROUP_NUM, BLOCK_GROUP_DIM] =>
some (Out, outOffset s DestLoc stride_o_bs stride_o_h stride_o_g stride_o_d idx))
(expected := fun idx : TileIndex [BLOCK_GROUP_NUM, BLOCK_GROUP_DIM] =>
if active s group_size BLOCK_GROUP_NUM BLOCK_GROUP_DIM idx then
quantizeKvCopyGroupSurfaceIntValue s K stride_k_bs stride_k_h
stride_k_g group_size BLOCK_GROUP_DIM hD idx
else s.readMemValue .int Out
(outOffset s DestLoc stride_o_bs stride_o_h stride_o_g stride_o_d idx)) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := destindex_copy_quantize_kv_group_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_g stride_k_d stride_o_bs stride_o_h
stride_o_g stride_o_d stride_os_bs stride_os_h stride_os_g group_size
BLOCK_GROUP_NUM BLOCK_GROUP_DIM)
(initialState := s)
(write := fun i : Fin BLOCK_GROUP_NUM =>
some (OutScale, scaleOutOffset s DestLoc stride_os_bs stride_os_h i))
(expected := fun i : Fin BLOCK_GROUP_NUM =>
if scaleActive group_size BLOCK_GROUP_NUM i then
quantizeKvCopyGroupScaleCell s K stride_k_bs stride_k_h stride_k_g
group_size BLOCK_GROUP_DIM hD i.val
else s.readMem OutScale
(scaleOutOffset s DestLoc stride_os_bs stride_os_h i))
Referenced model/specification definitions: Python source |
|---|
quantize_kv_copyquantize_kv_copy_io_correctnessMathematical execution | Configured model / stage Groupwise KV quantization and scale outputs in the mathematical conversion model. Source has repeated Python kernel names; source links identify definitions by line. No concrete int8 hardware guarantee. | Statement and preconditionsspecification quantize_kv_copy_io_correctness
(K : RegionName) (DestLoc : Region .nat) (Out OutScale : RegionName)
(stride_k_bs stride_k_h stride_k_g stride_k_d stride_o_bs stride_o_h stride_o_g stride_o_d
stride_os_bs stride_os_h stride_os_g
group_size BLOCK_GROUP_NUM BLOCK_GROUP_DIM : Nat)
(hD : 0 < BLOCK_GROUP_DIM) (hOut : Out ≠ OutScale)
(hValInj : Function.Injective
(fun idx : TileIndex [BLOCK_GROUP_NUM, BLOCK_GROUP_DIM] =>
idx.1.val * stride_o_g + idx.2.1.val)) :
quantizeKvCopyIO K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_g stride_k_d stride_o_bs stride_o_h stride_o_g
stride_o_d stride_os_bs stride_os_h stride_os_g group_size BLOCK_GROUP_NUM BLOCK_GROUP_DIM
⊨ fun _ _ _ xs =>
(fun j => valueXs BLOCK_GROUP_NUM BLOCK_GROUP_DIM group_size hD xs j,
fun i => scaleCellXs BLOCK_GROUP_NUM BLOCK_GROUP_DIM group_size hD xs i)
IO definition: quantizeKvCopyIOdef quantizeKvCopyIO
(K : RegionName) (DestLoc : Region .nat) (Out OutScale : RegionName)
(stride_k_bs stride_k_h stride_k_g stride_k_d stride_o_bs stride_o_h stride_o_g stride_o_d
stride_os_bs stride_os_h stride_os_g
group_size BLOCK_GROUP_NUM BLOCK_GROUP_DIM : Nat) :
MetaMasked2DKernelIO₁ₓ₂ where
kernel := destindex_copy_quantize_kv_group_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_g stride_k_d stride_o_bs stride_o_h stride_o_g
stride_o_d stride_os_bs stride_os_h stride_os_g group_size BLOCK_GROUP_NUM BLOCK_GROUP_DIM
mbuf1 := DestLoc
inp := K
out1 := Out
out2 := OutScale
B := BLOCK_GROUP_NUM * BLOCK_GROUP_DIM
C := BLOCK_GROUP_NUM
oty1 := .int
oty2 := .float
mwin1 := fun pid₀ _ => pid₀
read := fun pid₀ pid₁ _ j =>
pid₀ * stride_k_bs + pid₁ * stride_k_h +
(j.val / BLOCK_GROUP_DIM) * stride_k_g + (j.val % BLOCK_GROUP_DIM)
write1 := fun _ pid₁ m1 j =>
m1 * stride_o_bs + pid₁ * stride_o_h +
(j.val / BLOCK_GROUP_DIM) * stride_o_g + (j.val % BLOCK_GROUP_DIM)
write2 := fun _ pid₁ m1 i => m1 * stride_os_bs + pid₁ * stride_os_h + i.val
mask := fun _ _ _ j => j.val / BLOCK_GROUP_DIM < group_size
writeMask2 := fun _ _ _ i => i.val < group_size
Referenced model/specification definitions: Python source |
|---|
quantize_kv_transformdestindex_copy_quantize_kv_transform_output_summary_generalMathematical execution | Configured model / stage KV quantization/transform and scale outputs for the explicit head-dimension and destination layout. Mathematical casts do not certify fixed-width hardware quantization. | Statement and preconditionsspecification destindex_copy_quantize_kv_transform_output_summary_general
(K DestLoc Out OutScale : RegionName)
(stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num head_dim BLOCK_DMODEL BLOCK_HEAD : Nat)
(s : BlockState) (hD : 0 < BLOCK_DMODEL) (hOut : Out ≠ OutScale)
(hValInj : Function.Injective
(fun idx : TileIndex [BLOCK_HEAD, BLOCK_DMODEL] =>
outOffset s DestLoc stride_o_bs stride_o_h stride_o_d idx))
(hScaleInj : Function.Injective
(fun i : Fin BLOCK_HEAD => scaleOutOffset s DestLoc stride_os_bs stride_os_h i)) :
(∃ alg, (destindex_copy_quantize_kv_transform_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num head_dim BLOCK_DMODEL BLOCK_HEAD).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := destindex_copy_quantize_kv_transform_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num head_dim BLOCK_DMODEL BLOCK_HEAD)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_HEAD, BLOCK_DMODEL] =>
some (Out, outOffset s DestLoc stride_o_bs stride_o_h stride_o_d idx))
(expected := fun idx : TileIndex [BLOCK_HEAD, BLOCK_DMODEL] =>
(if active s head_num head_dim BLOCK_HEAD BLOCK_DMODEL idx then
quantizeKvTransformSurfaceIntValue s K stride_k_bs stride_k_h stride_k_d head_num head_dim BLOCK_DMODEL hD idx
else s.readMemValue .int Out (outOffset s DestLoc stride_o_bs stride_o_h stride_o_d idx))) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := destindex_copy_quantize_kv_transform_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num head_dim BLOCK_DMODEL BLOCK_HEAD)
(initialState := s)
(write := fun i : Fin BLOCK_HEAD =>
some (OutScale, scaleOutOffset s DestLoc stride_os_bs stride_os_h i))
(expected := fun i : Fin BLOCK_HEAD =>
(if scaleActive head_num BLOCK_HEAD i then
quantizeKvTransformScaleCell s K stride_k_bs stride_k_h stride_k_d head_num head_dim BLOCK_DMODEL hD i.val
else s.readMem OutScale (scaleOutOffset s DestLoc stride_os_bs stride_os_h i)))
Referenced model/specification definitions: Python source |
|---|
quantize_kv_transformquantize_kv_transform_io_correctnessMathematical execution | Configured model / stage KV quantization/transform and scale outputs for the explicit head-dimension and destination layout. Mathematical casts do not certify fixed-width hardware quantization. | Statement and preconditionsspecification quantize_kv_transform_io_correctness
(K : RegionName) (DestLoc : RegionName) (Out OutScale : RegionName)
(stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d
head_num head_dim BLOCK_DMODEL BLOCK_HEAD : Nat)
(hD : 0 < BLOCK_DMODEL) (hOut : Out ≠ OutScale)
(hValInj : Function.Injective
(fun idx : TileIndex [BLOCK_HEAD, BLOCK_DMODEL] =>
stride_o_h * idx.1.val + stride_o_d * idx.2.1.val))
(hScaleInj : Function.Injective
(fun i : Fin BLOCK_HEAD => stride_os_h * i.val)) :
quantizeKvTransformIO K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num head_dim BLOCK_DMODEL BLOCK_HEAD
⊨ fun _ _ _ xs =>
(fun j => valueXs BLOCK_HEAD BLOCK_DMODEL head_num head_dim hD xs j,
fun i => scaleCellXs BLOCK_HEAD BLOCK_DMODEL head_num head_dim hD xs i)
IO definition: quantizeKvTransformIOdef quantizeKvTransformIO
(K : RegionName) (DestLoc : RegionName) (Out OutScale : RegionName)
(stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d
head_num head_dim BLOCK_DMODEL BLOCK_HEAD : Nat) :
MetaMasked2DKernelIO₁ₓ₂ where
kernel := destindex_copy_quantize_kv_transform_real_surface K DestLoc Out OutScale
stride_k_bs stride_k_h stride_k_d stride_o_bs stride_o_h stride_o_d
stride_os_bs stride_os_h stride_os_d head_num head_dim BLOCK_DMODEL BLOCK_HEAD
mbuf1 := DestLoc
inp := K
out1 := Out
out2 := OutScale
B := BLOCK_HEAD * BLOCK_DMODEL
C := BLOCK_HEAD
oty1 := .int
oty2 := .float
mwin1 := fun pid₀ _ => pid₀
read := fun pid₀ _ _ j =>
pid₀ * stride_k_bs + (j.val / BLOCK_DMODEL) * stride_k_h +
stride_k_d * (j.val % BLOCK_DMODEL)
write1 := fun _ _ m1 j =>
m1 * stride_o_bs + stride_o_h * (j.val / BLOCK_DMODEL) +
stride_o_d * (j.val % BLOCK_DMODEL)
write2 := fun _ _ m1 i => m1 * stride_os_bs + stride_os_h * i.val
mask := fun _ _ _ j =>
j.val / BLOCK_DMODEL < head_num ∧ j.val % BLOCK_DMODEL < head_dim
writeMask2 := fun _ _ _ i => i.val < head_num
Referenced model/specification definitions: Python source |
|---|
rbe_triton_transformrbe_triton_transform_output_summary_generalMathematical execution | Configured model / stage Selected rotary pair outputs using mathematical sin/cos and the declared token-frequency model. Host frequency construction and concrete transcendental behavior are external. | Statement and preconditionsspecification rbe_triton_transform_output_summary_general
(x_ptr out_ptr : RegionName)
(M K stride_x_batch stride_x_m stride_x_n
stride_out_batch stride_out_m stride_out_n
start_token_position : Nat)
(THETA : ℝ) (DIM BLOCK_SIZE_M BLOCK_SIZE_K : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun idx : TileIndex [BLOCK_SIZE_M, BLOCK_SIZE_K / 2] =>
outOff s K stride_out_batch stride_out_m stride_out_n BLOCK_SIZE_M
BLOCK_SIZE_K idx))
(hRI : ∀ idx k : TileIndex [BLOCK_SIZE_M, BLOCK_SIZE_K / 2],
outOff s K stride_out_batch stride_out_m stride_out_n BLOCK_SIZE_M
BLOCK_SIZE_K idx
≠ outOff s K stride_out_batch stride_out_m stride_out_n BLOCK_SIZE_M
BLOCK_SIZE_K k + 1)
(hOutInjAll : ∀ s : BlockState, Function.Injective
(fun idx : TileIndex [BLOCK_SIZE_M, BLOCK_SIZE_K / 2] =>
outOff s K stride_out_batch stride_out_m stride_out_n BLOCK_SIZE_M
BLOCK_SIZE_K idx))
(hRIAll : ∀ (s : BlockState)
(idx k : TileIndex [BLOCK_SIZE_M, BLOCK_SIZE_K / 2]),
outOff s K stride_out_batch stride_out_m stride_out_n BLOCK_SIZE_M
BLOCK_SIZE_K idx
≠ outOff s K stride_out_batch stride_out_m stride_out_n BLOCK_SIZE_M
BLOCK_SIZE_K k + 1) :
-- (1) the full faithful surface lowers to the algorithm layer
(∃ alg, (rbe_triton_surface x_ptr out_ptr M K stride_x_batch stride_x_m
stride_x_n stride_out_batch stride_out_m stride_out_n
start_token_position THETA DIM BLOCK_SIZE_M BLOCK_SIZE_K).toAlgorithm?
= Except.ok alg) ∧
-- (2) even offsets: genuine `x_real·cos − x_imag·sin`
ComputeCorrect.Realizes_without_Rounding
(kernel := rbe_triton_surface x_ptr out_ptr M K stride_x_batch
stride_x_m stride_x_n stride_out_batch stride_out_m stride_out_n
start_token_position THETA DIM BLOCK_SIZE_M BLOCK_SIZE_K)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_SIZE_M, BLOCK_SIZE_K / 2] =>
activeReal s M K BLOCK_SIZE_M BLOCK_SIZE_K idx)
(fun idx => (out_ptr,
outOff s K stride_out_batch stride_out_m stride_out_n BLOCK_SIZE_M
BLOCK_SIZE_K idx)))
(expected := fun idx =>
rbeOutRealSpec s x_ptr K stride_x_batch stride_x_m stride_x_n
start_token_position DIM BLOCK_SIZE_M BLOCK_SIZE_K THETA idx) ∧
-- (3) odd offsets: genuine `x_real·sin + x_imag·cos`
ComputeCorrect.Realizes_without_Rounding
(kernel := rbe_triton_surface x_ptr out_ptr M K stride_x_batch
stride_x_m stride_x_n stride_out_batch stride_out_m stride_out_n
start_token_position THETA DIM BLOCK_SIZE_M BLOCK_SIZE_K)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_SIZE_M, BLOCK_SIZE_K / 2] =>
activeImag s M K BLOCK_SIZE_M BLOCK_SIZE_K idx)
(fun idx => (out_ptr,
outOff s K stride_out_batch stride_out_m stride_out_n BLOCK_SIZE_M
BLOCK_SIZE_K idx + 1)))
(expected := fun idx =>
rbeOutImagSpec s x_ptr K stride_x_batch stride_x_m stride_x_n
start_token_position DIM BLOCK_SIZE_M BLOCK_SIZE_K THETA idx) ∧
-- (4) the flat-memory `⊨` face of both interleaved stores
(rbeTritonIO x_ptr out_ptr M K stride_x_batch stride_x_m stride_x_n
stride_out_batch stride_out_m stride_out_n start_token_position THETA
DIM BLOCK_SIZE_M BLOCK_SIZE_K
⊨ fun _pid₀ pid₁ xs o j =>
let xr := xs (⟨0, by decide⟩ : Fin 2) j
let xi := xs (⟨1, by decide⟩ : Fin 2) j
let f := rbeFreqP pid₁ K start_token_position DIM BLOCK_SIZE_M
BLOCK_SIZE_K THETA j
match o with
| ⟨0, _⟩ =>
xr * Real.cos f -
(if 1 + rbeColP pid₁ K BLOCK_SIZE_M BLOCK_SIZE_K j < K then xi
else 0) * Real.sin f
| ⟨_ + 1, _⟩ => xr * Real.sin f + xi * Real.cos f)
IO definition: rbeTritonIOdef rbeTritonIO (x_ptr out_ptr : RegionName)
(M K stride_x_batch stride_x_m stride_x_n
stride_out_batch stride_out_m stride_out_n
start_token_position : Nat)
(THETA : ℝ) (DIM BLOCK_SIZE_M BLOCK_SIZE_K : Nat) :
GroupedMasked2DKernelIO where
kernel := rbe_triton_surface x_ptr out_ptr M K stride_x_batch stride_x_m
stride_x_n stride_out_batch stride_out_m stride_out_n
start_token_position THETA DIM BLOCK_SIZE_M BLOCK_SIZE_K
nIn := 2
nOut := 2
bufs := [x_ptr, out_ptr]
inp := fun _ => x_ptr
out := fun _ => out_ptr
B := BLOCK_SIZE_M * (BLOCK_SIZE_K / 2)
read := fun i pid₀ pid₁ j => match i with
| ⟨0, _⟩ => rbeXAddrP pid₀ pid₁ K stride_x_batch stride_x_m stride_x_n
BLOCK_SIZE_M BLOCK_SIZE_K j
| ⟨_ + 1, _⟩ => rbeXAddrP pid₀ pid₁ K stride_x_batch stride_x_m stride_x_n
BLOCK_SIZE_M BLOCK_SIZE_K j + 1
readMask := fun i _pid₀ pid₁ j => match i with
| ⟨0, _⟩ => rbeRowP pid₁ K BLOCK_SIZE_M BLOCK_SIZE_K j < M ∧
rbeColP pid₁ K BLOCK_SIZE_M BLOCK_SIZE_K j < K
| ⟨_ + 1, _⟩ => rbeRowP pid₁ K BLOCK_SIZE_M BLOCK_SIZE_K j < M ∧
1 + rbeColP pid₁ K BLOCK_SIZE_M BLOCK_SIZE_K j < K
write := fun o pid₀ pid₁ j => match o with
| ⟨0, _⟩ => rbeOutAddrP pid₀ pid₁ K stride_out_batch stride_out_m
stride_out_n BLOCK_SIZE_M BLOCK_SIZE_K j
| ⟨_ + 1, _⟩ => rbeOutAddrP pid₀ pid₁ K stride_out_batch stride_out_m
stride_out_n BLOCK_SIZE_M BLOCK_SIZE_K j + 1
writeMask := fun o _pid₀ pid₁ j => match o with
| ⟨0, _⟩ => rbeRowP pid₁ K BLOCK_SIZE_M BLOCK_SIZE_K j < M ∧
rbeColP pid₁ K BLOCK_SIZE_M BLOCK_SIZE_K j < K
| ⟨_ + 1, _⟩ => rbeRowP pid₁ K BLOCK_SIZE_M BLOCK_SIZE_K j < M ∧
1 + rbeColP pid₁ K BLOCK_SIZE_M BLOCK_SIZE_K j < K
Referenced model/specification definitions: Python source |
|---|
relu_strided_bufferrelu_strided_buffer_output_summary_generalMathematical execution | Configured model / stage Rank-one one-tile and grid-stride ReLU contracts. Positive output stride and, for streaming IO, input/output separation constrain applicability. | Statement and preconditionsspecification relu_strided_buffer_output_summary_general
(in0_ptr out0_ptr : RegionName)
(in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0 : Nat)
(s : BlockState)
(hStride : 0 < out0_stride0)
(hDisj : in0_ptr ≠ out0_ptr)
(hGrid : 0 < s.numPids 0) :
-- (1) both branch surfaces lower to the algorithm layer
(∃ alg, (relu_forward_kernel_rank_1_one_tile_surface in0_ptr out0_ptr
in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta
tile_size0).toAlgorithm? = Except.ok alg) ∧
(∃ alg, (relu_forward_kernel_rank_1_grid_stride_surface in0_ptr out0_ptr
in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta
tile_size0).toAlgorithm? = Except.ok alg) ∧
-- (2) one_tile_per_cta = true: genuine elementwise ReLU
ComputeCorrect.Realizes_without_Rounding
(kernel := relu_forward_kernel_rank_1_one_tile_surface in0_ptr out0_ptr
in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin tile_size0 => taskIndex (s.pids 0) tile_size0 i < s0)
(fun i => (out0_ptr, taskIndex (s.pids 0) tile_size0 i * out0_stride0)))
(expected := fun i =>
reluSpec s in0_ptr in0_stride0 (taskIndex (s.pids 0) tile_size0 i)) ∧
-- (3) one_tile_per_cta = false: genuine elementwise ReLU across the
-- whole grid-stride loop
ComputeCorrect.Realizes_without_Rounding
(kernel := relu_forward_kernel_rank_1_grid_stride_surface in0_ptr
out0_ptr in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun p : Fin tiles_per_cta × Fin tile_size0 =>
taskIndex (s.pids 0 + p.1.val * s.numPids 0) tile_size0 p.2 < s0)
(fun p => (out0_ptr,
taskIndex (s.pids 0 + p.1.val * s.numPids 0) tile_size0 p.2
* out0_stride0)))
(expected := fun p =>
reluSpec s in0_ptr in0_stride0
(taskIndex (s.pids 0 + p.1.val * s.numPids 0) tile_size0 p.2))
Referenced model/specification definitions: Python source |
|---|
relu_strided_bufferrelu_strided_buffer_one_tile_io_correctnessMathematical execution | Configured model / stage Rank-one one-tile and grid-stride ReLU contracts. Positive output stride and, for streaming IO, input/output separation constrain applicability. | Statement and preconditionsspecification relu_strided_buffer_one_tile_io_correctness
(in0_ptr out0_ptr : RegionName)
(in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0 : Nat)
(hStride : 0 < out0_stride0) :
reluOneTileIO in0_ptr out0_ptr in0_stride0 out0_stride0 s0 num_tasks
tiles_per_cta tile_size0
⊨ fun _pid xs i => TiledActivation.relu (xs i)
IO definition: reluOneTileIOdef reluOneTileIO (in0_ptr out0_ptr : RegionName)
(in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0 : Nat) :
MaskedTileKernelIO₁ where
kernel := relu_forward_kernel_rank_1_one_tile_surface in0_ptr out0_ptr
in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0
inp := in0_ptr
out := out0_ptr
shape := [tile_size0]
read := fun pid i => taskIndex pid tile_size0 i.1 * in0_stride0
write := fun pid i => taskIndex pid tile_size0 i.1 * out0_stride0
mask := fun pid i => taskIndex pid tile_size0 i.1 < s0
Referenced model/specification definitions: Python source |
|---|
relu_strided_bufferrelu_strided_buffer_grid_stride_io_correctnessRAbstract cast/store rounding | Configured model / stage Rank-one one-tile and grid-stride ReLU contracts. Positive output stride and, for streaming IO, input/output separation constrain applicability. | Statement and preconditionsspecification relu_strided_buffer_grid_stride_io_correctnessR (R : RoundingModel)
(in0_ptr out0_ptr : RegionName)
(in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0 : Nat)
(hStride : 0 < out0_stride0) (hDisj : in0_ptr ≠ out0_ptr) :
reluGridStrideIO in0_ptr out0_ptr in0_stride0 out0_stride0 s0 num_tasks
tiles_per_cta tile_size0
⊨[R] fun _pid₀ _pid₁ _nCtas xs t i => TiledActivation.relu (xs t i)
IO definition: reluGridStrideIOdef reluGridStrideIO (in0_ptr out0_ptr : RegionName)
(in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0 : Nat) :
StreamGridStrideEmitMasked2DKernelIO₁ where
kernel := relu_forward_kernel_rank_1_grid_stride_surface in0_ptr out0_ptr
in0_stride0 out0_stride0 s0 num_tasks tiles_per_cta tile_size0
inp1 := in0_ptr
out := out0_ptr
T := tiles_per_cta
B1 := tile_size0
C := tile_size0
pre := fun pid₀ _ nCtas => pid₀ < nCtas
read1 := fun p₀ _ nCtas t j =>
taskIndex (p₀ + t.val * nCtas) tile_size0 j * in0_stride0
write := fun p₀ _ nCtas t j =>
taskIndex (p₀ + t.val * nCtas) tile_size0 j * out0_stride0
mask1 := fun p₀ _ nCtas t j => taskIndex (p₀ + t.val * nCtas) tile_size0 j < s0
writeMask := fun p₀ _ nCtas t j =>
taskIndex (p₀ + t.val * nCtas) tile_size0 j < s0
Referenced model/specification definitions: Python source |
|---|
relu_triton_kernelrelu_kernel_correctnessMathematical execution | Configured model / stage Masked per-program ReLU with positive block size and the declared input/output windows. No host launch or concrete floating exceptional-value claim. | Statement and preconditionsspecification relu_kernel_correctness
(x_ptr out_ptr : RegionName) (N block_size : Nat)
(hB : 0 < block_size) :
reluIO x_ptr out_ptr N block_size ⊨
fun xs i => TiledActivation.relu (xs i)
IO definition: reluIOdef reluIO (x_ptr out_ptr : RegionName) (N block_size : Nat) :
MaskedKernelIO₁ where
kernel := relu_kernel x_ptr out_ptr N block_size
inp := x_ptr
out := out_ptr
B := block_size
read := fun pid => pid * block_size
write := fun pid => pid * block_size
mask := fun pid j => pid * block_size + j.val < N
writeMask := fun pid j => pid = 0 ∧ pid * block_size + j.val < N
Referenced model/specification definitions: Python source |
|---|
reversed_cumsumreversed_cumsum_output_summary_generalMathematical execution | Configured model / stage Reverse-direction cumulative-sum value facts for full/single-block surfaces under the named premises. | Statement and preconditionsspecification reversed_cumsum_output_summary_general
(SReg Z : RegionName) (s_s_h s_s_t s_s_d T S BT BS : Nat) (s : BlockState)
(hNoCollision : ∀ idx : TileIndex [BT, BS], singleBlockActive s T S BS idx →
∀ k : TileIndex [BT, BS], singleBlockActive s T S BS k →
singleBlockTileOffset s s_s_h s_s_t s_s_d BS k =
singleBlockTileOffset s s_s_h s_s_t s_s_d BS idx → k = idx) :
(∃ alg, (reversed_cumsum_surface SReg Z s_s_h s_s_t s_s_d T S BT BS).toAlgorithm? =
Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := reversed_cumsum_single_block_surface SReg Z s_s_h s_s_t s_s_d
T S BT BS)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BT, BS] => singleBlockActive s T S BS idx)
(fun idx : TileIndex [BT, BS] =>
(Z, singleBlockTileOffset s s_s_h s_s_t s_s_d BS idx)))
(expected := fun idx : TileIndex [BT, BS] =>
reversedCumsumClosed s SReg s_s_h s_s_t s_s_d T S BT BS idx))
Referenced model/specification definitions: Python source |
|---|
reversed_cumsumreversed_cumsum_block_store_io_correctnessMathematical execution | Precomputed-input slice Copies precomputed reverse-cumsum tile BC to Z; this IO slice does not compute the reverse sum. | Statement and preconditionsspecification reversed_cumsum_block_store_io_correctness (BC Z : RegionName)
(s_s_h s_s_t s_s_d T S BT BS : Nat)
(hOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BT, BS] =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t
+ (p₀ * BS + idx.2.1.val) * s_s_d)) :
blockStoreIO BC Z s_s_h s_s_t s_s_d T S BT BS
⊨ fun _p₀ _p₁ xs idx => xs idx
IO definition: blockStoreIOdef blockStoreIO (BC Z : RegionName) (s_s_h s_s_t s_s_d T S BT BS : Nat) :
Masked3DTileKernelIO₁ where
kernel := reversed_cumsum_store_slice BC Z s_s_h s_s_t s_s_d T S BT BS
inp := BC
out := Z
shape := [BT, BS]
read := fun p₀ p₁ p₂ idx =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t + (p₀ * BS + idx.2.1.val) * s_s_d
write := fun p₀ p₁ p₂ idx =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t + (p₀ * BS + idx.2.1.val) * s_s_d
mask := fun p₀ _p₁ p₂ idx =>
p₂ * BT + idx.1.val < T ∧ p₀ * BS + idx.2.1.val < S
Referenced model/specification definitions: Python source |
|---|
reversed_cumsumreversed_cumsum_block_store_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding contract for the supplied reverse-cumsum block-store slice. | Statement and preconditionsspecification reversed_cumsum_block_store_io_correctnessR (R : RoundingModel)
(BC Z : RegionName) (s_s_h s_s_t s_s_d T S BT BS : Nat)
(hOutInj : ∀ p₀ p₁ p₂ : Nat, Function.Injective
(fun idx : TileIndex [BT, BS] =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t
+ (p₀ * BS + idx.2.1.val) * s_s_d)) :
blockStoreIO BC Z s_s_h s_s_t s_s_d T S BT BS
⊨[R, FloatDType.real] fun _p₀ _p₁ xs idx => xs idx
IO definition: blockStoreIOdef blockStoreIO (BC Z : RegionName) (s_s_h s_s_t s_s_d T S BT BS : Nat) :
Masked3DTileKernelIO₁ where
kernel := reversed_cumsum_store_slice BC Z s_s_h s_s_t s_s_d T S BT BS
inp := BC
out := Z
shape := [BT, BS]
read := fun p₀ p₁ p₂ idx =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t + (p₀ * BS + idx.2.1.val) * s_s_d
write := fun p₀ p₁ p₂ idx =>
p₁ * s_s_h + (p₂ * BT + idx.1.val) * s_s_t + (p₀ * BS + idx.2.1.val) * s_s_d
mask := fun p₀ _p₁ p₂ idx =>
p₂ * BT + idx.1.val < T ∧ p₀ * BS + idx.2.1.val < S
Referenced model/specification definitions: Python source |
|---|
reversed_cumsum_scalarreversed_cumsum_scalar_output_summary_generalMathematical execution | Configured model / stage Scalar reverse-sum full/step facts with an explicit carry relation; summary conjuncts concern distinct surfaces. | Statement and preconditionsspecification reversed_cumsum_scalar_output_summary_general
(S Carry O : RegionName) (T BT : Nat) (hT : T ≤ BT) (s : BlockState)
(hOutInj : Function.Injective (fun i : Fin BT => singleBlockVecOffset s T i)) :
-- (1) the full reverse-traversal surface lowers to the algorithm layer
(∃ alg, (reversed_cumsum_scalar_surface S O T BT).toAlgorithm?
= Except.ok alg) ∧
-- (2) single-chunk genuine suffix sum, end-to-end from `S`
(ComputeCorrect.Realizes_without_Rounding
(kernel := reversed_cumsum_scalar_single_block_surface S O T BT)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BT => singleBlockActive s T i)
(fun i => (O, singleBlockVecOffset s T i)))
(expected := fun i : Fin BT => singleBlockRevClosed s S T BT i)) ∧
-- (3) per-chunk reversed carry-fold face, for every carry-invariant state
(∀ (s' : BlockState),
Function.Injective (fun i : Fin BT => vecOffset s' T BT i) →
s'.readMem Carry (s'.pids 0)
= (∑ flat ∈ (Finset.range T).filter
(fun flat => s'.pids 1 * BT ≤ flat), rowElem s' S T flat) →
ComputeCorrect.Realizes_without_Rounding
(kernel := reversed_cumsum_scalar_rev_slice S Carry O T BT)
(initialState := s')
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BT => active s' T BT i)
(fun i => (O, vecOffset s' T BT i)))
(expected := fun i : Fin BT => globalRevCumsumClosed s' S T BT i))
Referenced model/specification definitions: Python source |
|---|
reversed_cumsum_scalarreversed_cumsum_scalar_vec_store_io_correctnessMathematical execution | Precomputed-input slice Copies supplied BO values into O; carry construction and reverse summation are outside the IO slice. | Statement and preconditionsspecification reversed_cumsum_scalar_vec_store_io_correctness
(BO O : RegionName) (T BT : Nat)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun i : Fin BT => p₀ * T + (p₁ * BT + i.val))) :
vecStoreIO BO O T BT ⊨ fun _p₀ _p₁ xs i => xs i
IO definition: vecStoreIOdef vecStoreIO (BO O : RegionName) (T BT : Nat) : Masked3DTileKernelIO₁ where
kernel := reversed_cumsum_scalar_store_slice BO O T BT
inp := BO
out := O
shape := [BT]
read := fun p₀ p₁ _p₂ i => p₀ * T + (p₁ * BT + i.1.val)
write := fun p₀ p₁ _p₂ i => p₀ * T + (p₁ * BT + i.1.val)
mask := fun _p₀ p₁ _p₂ i => p₁ * BT + i.1.val < T
Referenced model/specification definitions: Python source |
|---|
reversed_cumsum_scalarreversed_cumsum_scalar_vec_store_io_correctnessRAbstract cast/store rounding | Precomputed-input slice Abstract rounding contract for the supplied BO-to-O store slice. | Statement and preconditionsspecification reversed_cumsum_scalar_vec_store_io_correctnessR
(R : RoundingModel) (BO O : RegionName) (T BT : Nat)
(hOutInj : ∀ p₀ p₁ : Nat, Function.Injective
(fun i : Fin BT => p₀ * T + (p₁ * BT + i.val))) :
vecStoreIO BO O T BT ⊨[R, FloatDType.real] fun _p₀ _p₁ xs i => xs i
IO definition: vecStoreIOdef vecStoreIO (BO O : RegionName) (T BT : Nat) : Masked3DTileKernelIO₁ where
kernel := reversed_cumsum_scalar_store_slice BO O T BT
inp := BO
out := O
shape := [BT]
read := fun p₀ p₁ _p₂ i => p₀ * T + (p₁ * BT + i.1.val)
write := fun p₀ p₁ _p₂ i => p₀ * T + (p₁ * BT + i.1.val)
mask := fun _p₀ p₁ _p₂ i => p₁ * BT + i.1.val < T
Referenced model/specification definitions: Python source |
|---|
rms_matmul_rberms_matmul_rbe_closed_form_correctMathematical execution | Configured model / stage Selected RMS/matmul and Q/K/V output surfaces at explicit reduction and stride configurations. Inspect flags and typed casts; the headline does not establish arbitrary rotary/host composition. | Statement and preconditionsspecification rms_matmul_rbe_closed_form_correct
(X W RMS OUT : RegionName) (s : BlockState)
(M N sxb sxm sxk swk swn srms sob som son BM BN BK numKBlocks : Nat) (K : Nat)
(hK : K = BK * numKBlocks) (EPS : ℝ)
(hcn : son = 1) (hbnle : BN ≤ som)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := rms_matmul_rbe_surface X W RMS OUT M N K sxb sxm sxk swk swn srms
sob som son BM BN BK numKBlocks EPS)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BM BN)
(fun idx => (OUT, outOffset s N BM BN sob som son idx)))
(expected := fun idx : TileIndex [BM, BN] =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (rmsSpec s X W RMS M N BM BN sxb sxm sxk swk swn srms
BK numKBlocks K EPS idx.1 idx.2.1))))
Referenced model/specification definitions: Python source |
|---|
rms_matmul_rberms_matmul_rbe_qkv_closed_form_correctMathematical execution | Configured model / stage Selected RMS/matmul and Q/K/V output surfaces at explicit reduction and stride configurations. Inspect flags and typed casts; the headline does not establish arbitrary rotary/host composition. | Statement and preconditionsspecification rms_matmul_rbe_qkv_closed_form_correct
(X QW KW VW RMS Q KOut V : RegionName) (s : BlockState)
(M N sxb sxm sxk sqwk sqwn skwk skwn svwk svwn srms
sqb sqm sqn skb skm skn svb svm svn BM BN BK numKBlocks : Nat) (K : Nat)
(hK : K = BK * numKBlocks) (EPS : ℝ)
(hqn : sqn = 1) (hqm : BN ≤ sqm)
(hkn : skn = 1) (hkm : BN ≤ skm)
(hvn : svn = 1) (hvm : BN ≤ svm)
(hXQ : X ≠ Q) (hRMSQ : RMS ≠ Q) (hKWQ : KW ≠ Q) (hVWQ : VW ≠ Q)
(hXK : X ≠ KOut) (hRMSK : RMS ≠ KOut) (hVWK : VW ≠ KOut)
(hQK : Q ≠ KOut) (hQV : Q ≠ V) (hKV : KOut ≠ V)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := rms_matmul_rbe_qkv_surface X QW KW VW RMS Q KOut V M N K sxb sxm sxk
sqwk sqwn skwk skwn svwk svwn srms sqb sqm sqn skb skm skn svb svm svn
BM BN BK numKBlocks EPS)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BM BN)
(fun idx => (Q, outOffset s N BM BN sqb sqm sqn idx)))
(expected := fun idx : TileIndex [BM, BN] =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (rmsSpec s X QW RMS M N BM BN sxb sxm sxk sqwk sqwn srms
BK numKBlocks K EPS idx.1 idx.2.1))))
∧ ComputeCorrect.Realizes_without_Rounding
(kernel := rms_matmul_rbe_qkv_surface X QW KW VW RMS Q KOut V M N K sxb sxm sxk
sqwk sqwn skwk skwn svwk svwn srms sqb sqm sqn skb skm skn svb svm svn
BM BN BK numKBlocks EPS)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BM BN)
(fun idx => (KOut, outOffset s N BM BN skb skm skn idx)))
(expected := fun idx : TileIndex [BM, BN] =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (rmsSpec s X KW RMS M N BM BN sxb sxm sxk skwk skwn srms
BK numKBlocks K EPS idx.1 idx.2.1))))
∧ ComputeCorrect.Realizes_without_Rounding
(kernel := rms_matmul_rbe_qkv_surface X QW KW VW RMS Q KOut V M N K sxb sxm sxk
sqwk sqwn skwk skwn svwk svwn srms sqb sqm sqn skb skm skn svb svm svn
BM BN BK numKBlocks EPS)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BM BN)
(fun idx => (V, outOffset s N BM BN svb svm svn idx)))
(expected := fun idx : TileIndex [BM, BN] =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (rmsSpec s X VW RMS M N BM BN sxb sxm sxk svwk svwn srms
BK numKBlocks K EPS idx.1 idx.2.1))))
Referenced model/specification definitions: Python source |
|---|
rms_norm_tritonrms_norm_kernel_correctnessMathematical execution | Configured model / stage Per-row RMS normalization under the declared mask, epsilon, and positive output stride. Mathematical sqrt/division and IO bounds define the contract. | Statement and preconditionsspecification rms_norm_kernel_correctness
(Y X W : RegionName)
(y_stride_r y_stride_c x_stride_r x_stride_c N BLOCK_SIZE : Nat)
(eps : ℝ) (hyc : 0 < y_stride_c) :
rmsNormIO Y X W y_stride_r y_stride_c x_stride_r x_stride_c N eps
BLOCK_SIZE ⊨
fun _ _ xs ws i => rmsNormSpec N BLOCK_SIZE eps xs ws i
IO definition: rmsNormIOdef rmsNormIO (Y X W : RegionName)
(y_stride_r y_stride_c x_stride_r x_stride_c N : Nat)
(eps : ℝ) (BLOCK_SIZE : Nat) : Masked2DKernelIO₂ where
kernel := rms_norm_kernel Y X W y_stride_r y_stride_c x_stride_r x_stride_c
N eps BLOCK_SIZE
in1 := X
in2 := W
out := Y
B := BLOCK_SIZE
read1 := fun pid _ j => pid * x_stride_r + j.val * x_stride_c
read2 := fun _ _ j => j.val
write := fun pid _ j => pid * y_stride_r + j.val * y_stride_c
mask := fun _ _ j => j.val < N
Referenced model/specification definitions: Python source |
|---|
rms_rbe_matmulrms_matmul_rbe_closed_form_correctMathematical execution | Configured model / stage Selected RMS/matmul output value contract. The separate Python rbe_triton kernel and its host composition are not proved by this headline. | Statement and preconditionsspecification rms_matmul_rbe_closed_form_correct
(X W RMS OUT : RegionName) (s : BlockState)
(M N sxb sxm sxk swk swn srms sob som son BM BN BK numKBlocks : Nat) (K : Nat)
(hK : K = BK * numKBlocks) (EPS : ℝ)
(hcn : son = 1) (hbnle : BN ≤ som)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := rms_matmul_rbe_surface X W RMS OUT M N K sxb sxm sxk swk swn srms
sob som son BM BN BK numKBlocks EPS)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BM BN)
(fun idx => (OUT, outOffset s N BM BN sob som son idx)))
(expected := fun idx : TileIndex [BM, BN] =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (rmsSpec s X W RMS M N BM BN sxb sxm sxk swk swn srms
BK numKBlocks K EPS idx.1 idx.2.1))))
Referenced model/specification definitions: Python source |
|---|
rmsnorm_fusedrms_norm_fwd_fused_output_summaryMathematical execution | Configured model / stage Single-row fused RMS contract; IO requires 0<N<=BLOCK_SIZE. This is a selected one-block configuration, not an unrestricted-width host implementation. | Statement and preconditionsspecification rms_norm_fwd_fused_output_summary
(X Y W : RegionName) (stride N BLOCK_SIZE : Nat) (eps : ℝ)
(s : BlockState)
(hNpos : 0 < N) (hNle : N ≤ BLOCK_SIZE)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => yOffset s stride i)) :
(∃ alg, (rms_norm_fwd_fused X Y W stride N BLOCK_SIZE eps).toAlgorithm? =
Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := rms_norm_fwd_fused X Y W stride N BLOCK_SIZE eps)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < N)
(fun i => (Y, yOffset s stride i)))
(expected := fun i => rmsnormSpec s X W stride N BLOCK_SIZE eps i)
Referenced model/specification definitions: Python source |
|---|
rmsnorm_fusedrms_norm_fwd_fused_io_correctnessAbstract cast/store rounding | Configured model / stage Single-row fused RMS contract; IO requires 0<N<=BLOCK_SIZE. This is a selected one-block configuration, not an unrestricted-width host implementation. | Statement and preconditionsspecification rms_norm_fwd_fused_io_correctness (R : RoundingModel)
(X Y W : RegionName) (stride N BLOCK_SIZE : Nat) (eps : ℝ)
(hNpos : 0 < N) (hNle : N ≤ BLOCK_SIZE) :
rmsnormFusedKernelIO X Y W stride N BLOCK_SIZE eps ⊨[R]
fun _ _ xs ws t j => rmsFusedStreamSpec N BLOCK_SIZE eps xs ws t j
IO definition: rmsnormFusedKernelIOdef rmsnormFusedKernelIO (X Y W : RegionName) (stride N BLOCK_SIZE : Nat)
(eps : ℝ) : StreamEmitMasked2DKernelIO₂ where
kernel := rms_norm_fwd_fused X Y W stride N BLOCK_SIZE eps
inp1 := X
inp2 := W
out := Y
T := 1
B1 := BLOCK_SIZE
B2 := BLOCK_SIZE
C := BLOCK_SIZE
outDType := .real
read1 := fun p₀ _ t j => p₀ * stride + (t.val * BLOCK_SIZE + j.val)
read2 := fun _ _ t j => t.val * BLOCK_SIZE + j.val
write := fun p₀ _ t j => p₀ * stride + (t.val * BLOCK_SIZE + j.val)
mask1 := fun _ _ t j => t.val * BLOCK_SIZE + j.val < N
mask2 := fun _ _ t j => t.val * BLOCK_SIZE + j.val < N
writeMask := fun _ _ t j => t.val * BLOCK_SIZE + j.val < N
Referenced model/specification definitions: Python source |
|---|
rmsnorm_fused_llamarms_norm_fwd_fused_llama_output_summaryMathematical execution | Configured model / stage Single-row fused LLaMA RMS contract with 0<N<=BLOCK_SIZE in the IO theorem and explicit typed/abstract output rounding; no concrete fp16 guarantee. | Statement and preconditionsspecification rms_norm_fwd_fused_llama_output_summary
(X Y W : RegionName) (stride N BLOCK_SIZE : Nat) (eps : ℝ)
(s : BlockState)
(hNpos : 0 < N) (hNle : N ≤ BLOCK_SIZE)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE => yOffset s stride i)) :
(∃ alg, (rms_norm_fwd_fused_llama X Y W stride N BLOCK_SIZE eps).toAlgorithm? =
Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := rms_norm_fwd_fused_llama X Y W stride N BLOCK_SIZE eps)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => i.val < N)
(fun i => (Y, yOffset s stride i)))
(expected := fun i =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (rmsnormSpec s X W stride N BLOCK_SIZE eps i))))
Referenced model/specification definitions: Python source |
|---|
rmsnorm_fused_llamarms_norm_fwd_fused_llama_io_correctnessAbstract cast/store rounding | Configured model / stage Single-row fused LLaMA RMS contract with 0<N<=BLOCK_SIZE in the IO theorem and explicit typed/abstract output rounding; no concrete fp16 guarantee. | Statement and preconditionsspecification rms_norm_fwd_fused_llama_io_correctness (R : RoundingModel)
(X Y W : RegionName) (stride N BLOCK_SIZE : Nat) (eps : ℝ)
(hNpos : 0 < N) (hNle : N ≤ BLOCK_SIZE) :
rmsnormFusedLlamaKernelIO X Y W stride N BLOCK_SIZE eps ⊨[R]
fun _ _ xs ws t j => rmsLlamaStreamSpec N BLOCK_SIZE eps xs ws t j
IO definition: rmsnormFusedLlamaKernelIOdef rmsnormFusedLlamaKernelIO (X Y W : RegionName) (stride N BLOCK_SIZE : Nat)
(eps : ℝ) : StreamEmitMasked2DKernelIO₂ where
kernel := rms_norm_fwd_fused_llama X Y W stride N BLOCK_SIZE eps
inp1 := X
inp2 := W
out := Y
T := 1
B1 := BLOCK_SIZE
B2 := BLOCK_SIZE
C := BLOCK_SIZE
outDType := .fp16
read1 := fun p₀ _ t j => p₀ * stride + (t.val * BLOCK_SIZE + j.val)
read2 := fun _ _ t j => t.val * BLOCK_SIZE + j.val
write := fun p₀ _ t j => p₀ * stride + (t.val * BLOCK_SIZE + j.val)
mask1 := fun _ _ t j => t.val * BLOCK_SIZE + j.val < N
mask2 := fun _ _ t j => t.val * BLOCK_SIZE + j.val < N
writeMask := fun _ _ t j => t.val * BLOCK_SIZE + j.val < N
Referenced model/specification definitions: Python source |
|---|
rmsnorm_implementationrmsnorm_implementation_output_summaryMathematical execution | Configured model / stage Full modeled row stream under positive block size, stride, and non-aliasing hypotheses. Host tensor/layout construction and hardware rounding are external. | Statement and preconditionsspecification rmsnorm_implementation_output_summary
(x_ptr rms_w_ptr out_ptr : RegionName)
(stride_x_batch stride_x_m stride_x_k stride_rms_w
stride_out_batch stride_out_m stride_out_k N_SIZE BLOCK_N_SIZE : Nat)
(eps : ℝ) (s : BlockState)
(hBlockPos : 0 < BLOCK_N_SIZE)
(hStrideOutKPos : 0 < stride_out_k)
(hXOutNe : x_ptr ≠ out_ptr)
(hWOutNe : rms_w_ptr ≠ out_ptr) :
(∃ alg, (rmsnorm_implementation x_ptr rms_w_ptr out_ptr
stride_x_batch stride_x_m stride_x_k stride_rms_w
stride_out_batch stride_out_m stride_out_k N_SIZE BLOCK_N_SIZE
eps).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := rmsnorm_implementation x_ptr rms_w_ptr out_ptr
stride_x_batch stride_x_m stride_x_k stride_rms_w
stride_out_batch stride_out_m stride_out_k N_SIZE BLOCK_N_SIZE eps)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun _ : Fin N_SIZE => True)
(fun i => (out_ptr,
outColOffset s stride_out_batch stride_out_m stride_out_k i.val)))
(expected := fun i =>
rmsnormWeightedYFullNSpec s x_ptr rms_w_ptr stride_x_batch stride_x_m
stride_x_k stride_rms_w N_SIZE BLOCK_N_SIZE eps i)
Referenced model/specification definitions: Python source |
|---|
rmsnorm_implementationrmsnorm_implementation_io_correctnessAbstract cast/store rounding | Configured model / stage Full modeled row stream under positive block size, stride, and non-aliasing hypotheses. Host tensor/layout construction and hardware rounding are external. | Statement and preconditionsspecification rmsnorm_implementation_io_correctness (R : RoundingModel)
(x_ptr rms_w_ptr out_ptr : RegionName)
(stride_x_batch stride_x_m stride_x_k stride_rms_w
stride_out_batch stride_out_m stride_out_k N_SIZE BLOCK_N_SIZE : Nat)
(eps : ℝ)
(hBlockPos : 0 < BLOCK_N_SIZE)
(hXOutNe : x_ptr ≠ out_ptr) (hWOutNe : rms_w_ptr ≠ out_ptr)
(hStrideOutKPos : 0 < stride_out_k) :
rmsnormImplementationKernelIO x_ptr rms_w_ptr out_ptr stride_x_batch
stride_x_m stride_x_k stride_rms_w stride_out_batch stride_out_m
stride_out_k N_SIZE BLOCK_N_SIZE eps ⊨[R]
fun _ _ xs ws t j =>
rmsImplStreamSpec N_SIZE BLOCK_N_SIZE eps xs ws t j
IO definition: rmsnormImplementationKernelIOdef rmsnormImplementationKernelIO (x w o : RegionName)
(sxb sxm sxk srw sob som sok N B : Nat) (eps : ℝ) :
StreamEmitMasked2DKernelIO₂ where
kernel := rmsnorm_implementation x w o sxb sxm sxk srw sob som sok N B eps
inp1 := x
inp2 := w
out := o
T := rmsNumSteps N B
B1 := B
B2 := B
C := B
read1 := fun p₀ p₁ t j => p₀ * sxb + p₁ * sxm + (t.val * B + j.val) * sxk
read2 := fun _ _ t j => (t.val * B + j.val) * srw
write := fun p₀ p₁ t j => p₀ * sob + p₁ * som + (t.val * B + j.val) * sok
mask1 := fun _ _ t j => t.val * B + j.val < N
mask2 := fun _ _ t j => t.val * B + j.val < N
writeMask := fun _ _ t j => t.val * B + j.val < N
Referenced model/specification definitions: Python source |
|---|
rmsnorm_tritonrmsnorm_full_output_summaryMathematical execution | Configured model / stage Full modeled RMS row stream with positive B, positive output stride, and output/input separation. The statement exposes the mathematical spec and abstract store rounding. | Statement and preconditionsspecification rmsnorm_full_output_summary
(x w o : RegionName) (sxb sxm sxk srw sob som sok N B : Nat) (eps : ℝ)
(s : BlockState) (hB : 0 < B) (hNpos : 0 < N)
(hox : o ≠ x) (how : o ≠ w)
(hsok : 0 < sok) :
(∃ alg, (VeriTile.Bench.TritonBenchG.RmsnormTriton.rmsnorm_triton x w o
sxb sxm sxk srw sob som sok N B eps).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := VeriTile.Bench.TritonBenchG.RmsnormTriton.rmsnorm_triton x w o
sxb sxm sxk srw sob som sok N B eps)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun _ : Fin N => True)
(fun k => (o, outOff s sob som sok k.val)))
(expected := fun k : Fin N => rmsSpecFull s x w sxb sxm sxk srw N eps k.val)
Referenced model/specification definitions: Python source |
|---|
rmsnorm_tritonrmsnorm_full_io_correctnessAbstract cast/store rounding | Configured model / stage Full modeled RMS row stream with positive B, positive output stride, and output/input separation. The statement exposes the mathematical spec and abstract store rounding. | Statement and preconditionsspecification rmsnorm_full_io_correctness (R : RoundingModel)
(x w o : RegionName) (sxb sxm sxk srw sob som sok N B : Nat) (eps : ℝ)
(hB : 0 < B) (hox : o ≠ x) (how : o ≠ w) (hsok : 0 < sok) :
rmsnormKernelIO x w o sxb sxm sxk srw sob som sok N B eps ⊨[R]
fun _ _ xs ws t j => rmsStreamSpec N B eps xs ws t j
IO definition: rmsnormKernelIOdef rmsnormKernelIO (x w o : RegionName)
(sxb sxm sxk srw sob som sok N B : Nat) (eps : ℝ) :
StreamEmitMasked2DKernelIO₂ where
kernel := rmsnorm_triton x w o sxb sxm sxk srw sob som sok N B eps
inp1 := x
inp2 := w
out := o
T := rmsNumSteps N B
B1 := B
B2 := B
C := B
read1 := fun p₀ p₁ t j => p₀ * sxb + p₁ * sxm + (t.val * B + j.val) * sxk
read2 := fun _ _ t j => (t.val * B + j.val) * srw
write := fun p₀ p₁ t j => p₀ * sob + p₁ * som + (t.val * B + j.val) * sok
mask1 := fun _ _ t j => t.val * B + j.val < N
mask2 := fun _ _ t j => t.val * B + j.val < N
writeMask := fun _ _ t j => t.val * B + j.val < N
Referenced model/specification definitions: Python source |
|---|
rope_backward_transformrope_backward_python_backward_output_summary_generalMathematical execution | Configured model / stage Backward full-surface Q/K value facts plus separate head-store IO slices under region/half-window assumptions. The slice frames do not automatically compose into a full-surface IO frame. | Statement and preconditionsspecification rope_backward_python_backward_output_summary_general
(Q K COS SIN : RegionName)
(q_row_stride k_row_stride cos_row_stride sin_row_stride
sl bs n_qh n_kh hd pad_n_qh pad_n_kh pad_hd BLOCK_SIZE : Nat)
(HEAD_IDX COS_ROW_IDX HEAD_HALF BLOCK_HALF : Nat)
(s : BlockState) (hundef : ∀ rg o, s.undef rg o = 0) (hqk : Q ≠ K) :
ComputeCorrect.Realizes_without_Rounding
(kernel := triton_rope_surface Q K COS SIN q_row_stride k_row_stride
cos_row_stride sin_row_stride sl bs n_qh n_kh hd pad_n_qh pad_n_kh
pad_hd BLOCK_SIZE Bool.true)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [pad_n_qh, pad_hd/2] =>
activeQFull (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2) n_qh hd idx)
(fun idx => (Q, qFullFirstOffset (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2) s q_row_stride hd idx)))
(expected := fun idx =>
ropeBackwardKernelQ0Spec (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2) s Q COS SIN q_row_stride sl cos_row_stride sin_row_stride hd idx) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := triton_rope_surface Q K COS SIN q_row_stride k_row_stride
cos_row_stride sin_row_stride sl bs n_qh n_kh hd pad_n_qh pad_n_kh
pad_hd BLOCK_SIZE Bool.true)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [pad_n_qh, pad_hd/2] =>
activeQFull (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2) n_qh hd idx)
(fun idx => (Q, qFullSecondOffset (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2) s q_row_stride hd idx)))
(expected := fun idx =>
ropeBackwardKernelQ1Spec (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2) s Q COS SIN q_row_stride sl cos_row_stride sin_row_stride hd idx) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := triton_rope_surface Q K COS SIN q_row_stride k_row_stride
cos_row_stride sin_row_stride sl bs n_qh n_kh hd pad_n_qh pad_n_kh
pad_hd BLOCK_SIZE Bool.true)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [pad_n_kh, pad_hd/2] =>
activeKFull (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2) n_kh hd idx)
(fun idx => (K, kFullFirstOffset (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2) s k_row_stride hd idx)))
(expected := fun idx =>
ropeBackwardKernelK0Spec (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2) s K COS SIN k_row_stride sl cos_row_stride sin_row_stride hd idx) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := triton_rope_surface Q K COS SIN q_row_stride k_row_stride
cos_row_stride sin_row_stride sl bs n_qh n_kh hd pad_n_qh pad_n_kh
pad_hd BLOCK_SIZE Bool.true)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [pad_n_kh, pad_hd/2] =>
activeKFull (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2) n_kh hd idx)
(fun idx => (K, kFullSecondOffset (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2) s k_row_stride hd idx)))
(expected := fun idx =>
ropeBackwardKernelK1Spec (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2) s K COS SIN k_row_stride sl cos_row_stride sin_row_stride hd idx) ∧
(ropeBackwardQ0HeadIO Q COS SIN HEAD_IDX COS_ROW_IDX q_row_stride
cos_row_stride sin_row_stride hd n_qh HEAD_HALF BLOCK_HALF
⊨ fun _ _ xs _ j =>
let q0 := xs (⟨0, by decide⟩ : Fin 4) j
let q1 := xs (⟨1, by decide⟩ : Fin 4) j
let c := xs (⟨2, by decide⟩ : Fin 4) j
let sn := xs (⟨3, by decide⟩ : Fin 4) j
q0 * c + q1 * sn) ∧
(ropeBackwardQ1HeadIO Q COS SIN HEAD_IDX COS_ROW_IDX q_row_stride
cos_row_stride sin_row_stride hd n_qh HEAD_HALF BLOCK_HALF
⊨ fun _ _ xs _ j =>
let q0 := xs (⟨0, by decide⟩ : Fin 4) j
let q1 := xs (⟨1, by decide⟩ : Fin 4) j
let c := xs (⟨2, by decide⟩ : Fin 4) j
let sn := xs (⟨3, by decide⟩ : Fin 4) j
q1 * c - q0 * sn) ∧
(ropeBackwardK0HeadIO K COS SIN HEAD_IDX COS_ROW_IDX k_row_stride
cos_row_stride sin_row_stride hd n_kh HEAD_HALF BLOCK_HALF
⊨ fun _ _ xs _ j =>
let k0 := xs (⟨0, by decide⟩ : Fin 4) j
let k1 := xs (⟨1, by decide⟩ : Fin 4) j
let c := xs (⟨2, by decide⟩ : Fin 4) j
let sn := xs (⟨3, by decide⟩ : Fin 4) j
k0 * c + k1 * sn) ∧
(ropeBackwardK1HeadIO K COS SIN HEAD_IDX COS_ROW_IDX k_row_stride
cos_row_stride sin_row_stride hd n_kh HEAD_HALF BLOCK_HALF
⊨ fun _ _ xs _ j =>
let k0 := xs (⟨0, by decide⟩ : Fin 4) j
let k1 := xs (⟨1, by decide⟩ : Fin 4) j
let c := xs (⟨2, by decide⟩ : Fin 4) j
let sn := xs (⟨3, by decide⟩ : Fin 4) j
k1 * c - k0 * sn)
Referenced model/specification definitions: Python source |
|---|
rope_embeddingrope_embedding_forward_backward_summary_generalMathematical execution | Configured model / stage Four forward/backward half-store contracts plus projection facts for the full surface. Independent half-store proofs do not establish combined in-place execution without composition/disjointness evidence. | Statement and preconditionsspecification rope_embedding_forward_backward_summary_general
(Q dY cos sin : RegionName)
(Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE : Nat)
(sQ sDY : BlockState)
(hQF : Function.Injective
(fun i : Fin BLOCK_SIZE => qFirstOffset sQ Q_row_stride head_dim ROPE_GROUP_SIZE i))
(hQS : Function.Injective
(fun i : Fin BLOCK_SIZE => qSecondOffset sQ Q_row_stride head_dim ROPE_GROUP_SIZE i))
(hDF : Function.Injective
(fun i : Fin BLOCK_SIZE => qFirstOffset sDY Q_row_stride head_dim ROPE_GROUP_SIZE i))
(hDS : Function.Injective
(fun i : Fin BLOCK_SIZE => qSecondOffset sDY Q_row_stride head_dim ROPE_GROUP_SIZE i)) :
(∃ alg, (rope_embedding_surface Q Q_row_stride cos cos_row_stride sin
sin_row_stride seqlen head_dim n_heads Bool.false BLOCK_SIZE ROPE_GROUP_SIZE).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := rope_embedding_forward_first_half Q cos sin Q_row_stride
cos_row_stride sin_row_stride seqlen head_dim n_heads ROPE_GROUP_SIZE BLOCK_SIZE)
(initialState := sQ)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => active sQ head_dim n_heads ROPE_GROUP_SIZE BLOCK_SIZE i)
(fun i => (Q, qFirstOffset sQ Q_row_stride head_dim ROPE_GROUP_SIZE i)))
(expected := fun i => ropeFirstSpec sQ Q cos sin Q_row_stride cos_row_stride
sin_row_stride seqlen head_dim ROPE_GROUP_SIZE BLOCK_SIZE i) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := rope_embedding_forward_second_half Q cos sin Q_row_stride
cos_row_stride sin_row_stride seqlen head_dim n_heads ROPE_GROUP_SIZE BLOCK_SIZE)
(initialState := sQ)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => active sQ head_dim n_heads ROPE_GROUP_SIZE BLOCK_SIZE i)
(fun i => (Q, qSecondOffset sQ Q_row_stride head_dim ROPE_GROUP_SIZE i)))
(expected := fun i => ropeSecondSpec sQ Q cos sin Q_row_stride cos_row_stride
sin_row_stride seqlen head_dim ROPE_GROUP_SIZE BLOCK_SIZE i) ∧
(∃ alg, (rope_embedding_surface dY Q_row_stride cos cos_row_stride sin
sin_row_stride seqlen head_dim n_heads Bool.true BLOCK_SIZE ROPE_GROUP_SIZE).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := rope_embedding_backward_first_half dY cos sin Q_row_stride
cos_row_stride sin_row_stride seqlen head_dim n_heads ROPE_GROUP_SIZE BLOCK_SIZE)
(initialState := sDY)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => active sDY head_dim n_heads ROPE_GROUP_SIZE BLOCK_SIZE i)
(fun i => (dY, qFirstOffset sDY Q_row_stride head_dim ROPE_GROUP_SIZE i)))
(expected := fun i => ropeBackwardFirstSpec sDY dY cos sin Q_row_stride
cos_row_stride sin_row_stride seqlen head_dim ROPE_GROUP_SIZE BLOCK_SIZE i) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := rope_embedding_backward_second_half dY cos sin Q_row_stride
cos_row_stride sin_row_stride seqlen head_dim n_heads ROPE_GROUP_SIZE BLOCK_SIZE)
(initialState := sDY)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE => active sDY head_dim n_heads ROPE_GROUP_SIZE BLOCK_SIZE i)
(fun i => (dY, qSecondOffset sDY Q_row_stride head_dim ROPE_GROUP_SIZE i)))
(expected := fun i => ropeBackwardSecondSpec sDY dY cos sin Q_row_stride
cos_row_stride sin_row_stride seqlen head_dim ROPE_GROUP_SIZE BLOCK_SIZE i)
Referenced model/specification definitions: Python source |
|---|
rope_embeddingrope_embedding_io_correctnessMathematical execution | Configured model / stage Four forward/backward half-store contracts plus projection facts for the full surface. Independent half-store proofs do not establish combined in-place execution without composition/disjointness evidence. | Statement and preconditionsspecification rope_embedding_io_correctness (Q cos sin : RegionName)
(Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE : Nat) :
(fwd_firstIO Q cos sin Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE
⊨ fun _p₀ _p₁ q1 q2 c1 s1 i => ropeFirstSpecOf q1 q2 c1 s1 i) ∧
(fwd_secondIO Q cos sin Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE
⊨ fun _p₀ _p₁ q1 q2 c1 s1 i => ropeSecondSpecOf q1 q2 c1 s1 i) ∧
(bwd_firstIO Q cos sin Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE
⊨ fun _p₀ _p₁ q1 q2 c1 s1 i => ropeBackwardFirstSpecOf q1 q2 c1 s1 i) ∧
(bwd_secondIO Q cos sin Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE
⊨ fun _p₀ _p₁ q1 q2 c1 s1 i => ropeBackwardSecondSpecOf q1 q2 c1 s1 i)
IO definition: fwd_firstIOdef fwd_firstIO (Q cos sin : RegionName)
(Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE : Nat) : InPlaceMaskedTileKernelIO where
kernel := rope_embedding_forward_first_half Q cos sin Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE
main := Q
aux1 := cos
aux2 := sin
shape := [BLOCK_SIZE]
readMain1 := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val
readMain2 := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val + head_dim / 2
readAux1 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * cos_row_stride + i.1.val
readAux2 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * sin_row_stride + i.1.val
write := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val
mask := fun _p₀ _p₁ i => i.1.val < head_dim / 2
writeMask := fun _p₀ p₁ i =>
i.1.val < head_dim / 2 ∧ p₁ * ROPE_GROUP_SIZE < n_heads
IO definition: fwd_secondIOdef fwd_secondIO (Q cos sin : RegionName)
(Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE : Nat) : InPlaceMaskedTileKernelIO where
kernel := rope_embedding_forward_second_half Q cos sin Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE
main := Q
aux1 := cos
aux2 := sin
shape := [BLOCK_SIZE]
readMain1 := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val
readMain2 := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val + head_dim / 2
readAux1 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * cos_row_stride + i.1.val
readAux2 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * sin_row_stride + i.1.val
write := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val + head_dim / 2
mask := fun _p₀ _p₁ i => i.1.val < head_dim / 2
writeMask := fun _p₀ p₁ i =>
i.1.val < head_dim / 2 ∧ p₁ * ROPE_GROUP_SIZE < n_heads
IO definition: bwd_firstIOdef bwd_firstIO (Q cos sin : RegionName)
(Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE : Nat) : InPlaceMaskedTileKernelIO where
kernel := rope_embedding_backward_first_half Q cos sin Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE
main := Q
aux1 := cos
aux2 := sin
shape := [BLOCK_SIZE]
readMain1 := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val
readMain2 := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val + head_dim / 2
readAux1 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * cos_row_stride + i.1.val
readAux2 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * sin_row_stride + i.1.val
write := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val
mask := fun _p₀ _p₁ i => i.1.val < head_dim / 2
writeMask := fun _p₀ p₁ i =>
i.1.val < head_dim / 2 ∧ p₁ * ROPE_GROUP_SIZE < n_heads
IO definition: bwd_secondIOdef bwd_secondIO (Q cos sin : RegionName)
(Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE : Nat) : InPlaceMaskedTileKernelIO where
kernel := rope_embedding_backward_second_half Q cos sin Q_row_stride cos_row_stride sin_row_stride seqlen head_dim n_heads
ROPE_GROUP_SIZE BLOCK_SIZE
main := Q
aux1 := cos
aux2 := sin
shape := [BLOCK_SIZE]
readMain1 := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val
readMain2 := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val + head_dim / 2
readAux1 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * cos_row_stride + i.1.val
readAux2 := fun p₀ _p₁ i =>
IntegralDType.nat.mod p₀ seqlen * sin_row_stride + i.1.val
write := fun p₀ p₁ i =>
p₀ * Q_row_stride + p₁ * ROPE_GROUP_SIZE * head_dim + i.1.val + head_dim / 2
mask := fun _p₀ _p₁ i => i.1.val < head_dim / 2
writeMask := fun _p₀ p₁ i =>
i.1.val < head_dim / 2 ∧ p₁ * ROPE_GROUP_SIZE < n_heads
Referenced model/specification definitions: Python source |
|---|
rope_transformrope_transform_output_summary_generalMathematical execution | Configured model / stage Forward full-surface Q/K value facts plus separate head-store IO contracts. The top bundle distinguishes the full value result from individual slice frames and assumes distinct Q/K buffers. | Statement and preconditionsspecification rope_transform_output_summary_general
(Q K COS SIN : RegionName)
(q_row_stride k_row_stride cos_row_stride sin_row_stride
sl bs n_qh n_kh hd pad_n_qh pad_n_kh pad_hd BLOCK_SIZE : Nat)
(HEAD_IDX COS_ROW_IDX HEAD_HALF BLOCK_HALF : Nat)
(s : BlockState) (hundef : ∀ rg o, s.undef rg o = 0) (hqk : Q ≠ K) :
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_rope_surface Q K COS SIN q_row_stride k_row_stride
cos_row_stride sin_row_stride sl bs n_qh n_kh hd pad_n_qh pad_n_kh
pad_hd BLOCK_SIZE Bool.false)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [pad_n_qh, pad_hd/2] =>
activeQFull (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2) n_qh hd idx)
(fun idx => (Q,
qFullFirstOffset (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2) s q_row_stride hd idx)))
(expected := fun idx =>
ropeForwardKernelQ0Spec (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2)
s Q COS SIN q_row_stride sl cos_row_stride sin_row_stride hd idx)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_rope_surface Q K COS SIN q_row_stride k_row_stride
cos_row_stride sin_row_stride sl bs n_qh n_kh hd pad_n_qh pad_n_kh
pad_hd BLOCK_SIZE Bool.false)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [pad_n_qh, pad_hd/2] =>
activeQFull (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2) n_qh hd idx)
(fun idx => (Q,
qFullSecondOffset (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2) s q_row_stride hd idx)))
(expected := fun idx =>
ropeForwardKernelQ1Spec (pad_n_qh := pad_n_qh) (pad_hd_half := pad_hd/2)
s Q COS SIN q_row_stride sl cos_row_stride sin_row_stride hd idx)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_rope_surface Q K COS SIN q_row_stride k_row_stride
cos_row_stride sin_row_stride sl bs n_qh n_kh hd pad_n_qh pad_n_kh
pad_hd BLOCK_SIZE Bool.false)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [pad_n_kh, pad_hd/2] =>
activeKFull (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2) n_kh hd idx)
(fun idx => (K,
kFullFirstOffset (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2) s k_row_stride hd idx)))
(expected := fun idx =>
ropeForwardKernelK0Spec (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2)
s K COS SIN k_row_stride sl cos_row_stride sin_row_stride hd idx)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_rope_surface Q K COS SIN q_row_stride k_row_stride
cos_row_stride sin_row_stride sl bs n_qh n_kh hd pad_n_qh pad_n_kh
pad_hd BLOCK_SIZE Bool.false)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [pad_n_kh, pad_hd/2] =>
activeKFull (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2) n_kh hd idx)
(fun idx => (K,
kFullSecondOffset (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2) s k_row_stride hd idx)))
(expected := fun idx =>
ropeForwardKernelK1Spec (pad_n_kh := pad_n_kh) (pad_hd_half := pad_hd/2)
s K COS SIN k_row_stride sl cos_row_stride sin_row_stride hd idx)) ∧
(ropeQ0HeadIO Q COS SIN HEAD_IDX COS_ROW_IDX q_row_stride cos_row_stride
sin_row_stride hd n_qh HEAD_HALF BLOCK_HALF
⊨ fun _ _ xs _ j =>
let q0 := xs (⟨0, by decide⟩ : Fin 4) j
let q1 := xs (⟨1, by decide⟩ : Fin 4) j
let c := xs (⟨2, by decide⟩ : Fin 4) j
let sn := xs (⟨3, by decide⟩ : Fin 4) j
q0 * c - q1 * sn) ∧
(ropeQ1HeadIO Q COS SIN HEAD_IDX COS_ROW_IDX q_row_stride cos_row_stride
sin_row_stride hd n_qh HEAD_HALF BLOCK_HALF
⊨ fun _ _ xs _ j =>
let q0 := xs (⟨0, by decide⟩ : Fin 4) j
let q1 := xs (⟨1, by decide⟩ : Fin 4) j
let c := xs (⟨2, by decide⟩ : Fin 4) j
let sn := xs (⟨3, by decide⟩ : Fin 4) j
q1 * c + q0 * sn) ∧
(ropeK0HeadIO K COS SIN HEAD_IDX COS_ROW_IDX k_row_stride cos_row_stride
sin_row_stride hd n_kh HEAD_HALF BLOCK_HALF
⊨ fun _ _ xs _ j =>
let k0 := xs (⟨0, by decide⟩ : Fin 4) j
let k1 := xs (⟨1, by decide⟩ : Fin 4) j
let c := xs (⟨2, by decide⟩ : Fin 4) j
let sn := xs (⟨3, by decide⟩ : Fin 4) j
k0 * c - k1 * sn) ∧
(ropeK1HeadIO K COS SIN HEAD_IDX COS_ROW_IDX k_row_stride cos_row_stride
sin_row_stride hd n_kh HEAD_HALF BLOCK_HALF
⊨ fun _ _ xs _ j =>
let k0 := xs (⟨0, by decide⟩ : Fin 4) j
let k1 := xs (⟨1, by decide⟩ : Fin 4) j
let c := xs (⟨2, by decide⟩ : Fin 4) j
let sn := xs (⟨3, by decide⟩ : Fin 4) j
k1 * c + k0 * sn)
Referenced model/specification definitions: Python source |
|---|
rotary_embrotary_emb_kernel_correctnessMathematical execution | Configured model / stage Four independent q0/q1/k0/k1 block-store IO contracts and full-surface projection. The source explicitly does not claim the simultaneous four-store value/frame composition. | Statement and preconditionsspecification rotary_emb_kernel_correctness
(Q K Cos Sin : RegionName)
(stride_qbs stride_qh stride_qd stride_kbs stride_kh stride_kd
stride_cosbs stride_cosd stride_sinbs stride_sind max_total_len
HEAD_Q HEAD_K BLOCK_HEAD BLOCK_SEQ BLOCK_DMODEL BLOCK_HALF : Nat)
(hQEven : Function.Injective
(fun i : Fin BLOCK_HALF => dimEven i * stride_qd))
(hQOdd : Function.Injective
(fun i : Fin BLOCK_HALF => dimOdd i * stride_qd))
(hKEven : Function.Injective
(fun i : Fin BLOCK_HALF => dimEven i * stride_kd))
(hKOdd : Function.Injective
(fun i : Fin BLOCK_HALF => dimOdd i * stride_kd)) :
(∃ alg, (rotary_kernel_surface Q K Cos Sin stride_qbs stride_qh stride_qd
stride_kbs stride_kh stride_kd stride_cosbs stride_cosd stride_sinbs
stride_sind max_total_len HEAD_Q HEAD_K BLOCK_HEAD BLOCK_SEQ
BLOCK_DMODEL).toAlgorithm? = Except.ok alg) ∧
(rotaryQ0IO Q Cos Sin stride_qbs stride_qh stride_qd stride_cosbs
stride_cosd stride_sinbs stride_sind max_total_len HEAD_Q BLOCK_HALF
⊨ fun _ _ xs _ j =>
let dataEven := xs (⟨0, by decide⟩ : Fin 4) j
let dataOdd := xs (⟨1, by decide⟩ : Fin 4) j
let cosLane := xs (⟨2, by decide⟩ : Fin 4) j
let sinLane := xs (⟨3, by decide⟩ : Fin 4) j
dataEven * cosLane - dataOdd * sinLane) ∧
(rotaryQ1IO Q Cos Sin stride_qbs stride_qh stride_qd stride_cosbs
stride_cosd stride_sinbs stride_sind max_total_len HEAD_Q BLOCK_HALF
⊨ fun _ _ xs _ j =>
let dataEven := xs (⟨0, by decide⟩ : Fin 4) j
let dataOdd := xs (⟨1, by decide⟩ : Fin 4) j
let cosLane := xs (⟨2, by decide⟩ : Fin 4) j
let sinLane := xs (⟨3, by decide⟩ : Fin 4) j
dataEven * sinLane + dataOdd * cosLane) ∧
(rotaryK0IO K Cos Sin stride_kbs stride_kh stride_kd stride_cosbs
stride_cosd stride_sinbs stride_sind max_total_len HEAD_K BLOCK_HALF
⊨ fun _ _ xs _ j =>
let dataEven := xs (⟨0, by decide⟩ : Fin 4) j
let dataOdd := xs (⟨1, by decide⟩ : Fin 4) j
let cosLane := xs (⟨2, by decide⟩ : Fin 4) j
let sinLane := xs (⟨3, by decide⟩ : Fin 4) j
dataEven * cosLane - dataOdd * sinLane) ∧
(rotaryK1IO K Cos Sin stride_kbs stride_kh stride_kd stride_cosbs
stride_cosd stride_sinbs stride_sind max_total_len HEAD_K BLOCK_HALF
⊨ fun _ _ xs _ j =>
let dataEven := xs (⟨0, by decide⟩ : Fin 4) j
let dataOdd := xs (⟨1, by decide⟩ : Fin 4) j
let cosLane := xs (⟨2, by decide⟩ : Fin 4) j
let sinLane := xs (⟨3, by decide⟩ : Fin 4) j
dataEven * sinLane + dataOdd * cosLane)
Referenced model/specification definitions: Python source |
|---|
rotary_emb_nopadrotary_emb_nopad_output_summary_generalMathematical execution | Configured model / stage Selected rotary Q and cache sub-surface contracts with explicit head-half/metadata/layout hypotheses. They do not establish arbitrary aliasing or host composition of all fused stages. | Statement and preconditionsspecification rotary_emb_nopad_output_summary_general
(Q K Cos Sin KVCache : RegionName)
(BlockTables ContextLengths : Region .nat)
(surf_q_token_stride surf_q_head_stride surf_k_token_stride surf_k_head_stride
surf_head_dim_stride surf_cos_token_stride surf_cos_stride
surf_q_total_tokens surf_Q_HEAD_NUM surf_KV_GROUP_NUM surf_HEAD_DIM
surf_BLOCK_TOKENS : Nat)
(q_token_stride q_head_stride head_dim_stride cos_token_stride cos_stride
q_total_tokens Q_HEAD_NUM HEAD_HALF BLOCK_TOKENS : Nat)
(k_token_stride k_head_stride k_q_total_tokens KV_GROUP_NUM : Nat)
(hQInjFirst : ∀ s : BlockState, Function.Injective
(fun idx : TileIndex [BLOCK_TOKENS, 1, HEAD_HALF] =>
qFullFirstOffset s q_token_stride q_head_stride head_dim_stride
BLOCK_TOKENS idx))
(hQInjSecond : ∀ s : BlockState, Function.Injective
(fun idx : TileIndex [BLOCK_TOKENS, 1, HEAD_HALF] =>
qFullSecondOffset s q_token_stride q_head_stride head_dim_stride
BLOCK_TOKENS HEAD_HALF idx))
(hQDisjoint : ∀ (s : BlockState)
(idx idx' : TileIndex [BLOCK_TOKENS, 1, HEAD_HALF]),
qFullFirstOffset s q_token_stride q_head_stride head_dim_stride
BLOCK_TOKENS idx ≠
qFullSecondOffset s q_token_stride q_head_stride head_dim_stride
BLOCK_TOKENS HEAD_HALF idx')
(hKInjFirst : ∀ s : BlockState, Function.Injective
(fun idx : TileIndex [BLOCK_TOKENS, 1, HEAD_HALF] =>
kFullFirstOffset s k_token_stride k_head_stride head_dim_stride
KV_GROUP_NUM BLOCK_TOKENS idx))
(hKInjSecond : ∀ s : BlockState, Function.Injective
(fun idx : TileIndex [BLOCK_TOKENS, 1, HEAD_HALF] =>
kFullSecondOffset s k_token_stride k_head_stride head_dim_stride
KV_GROUP_NUM BLOCK_TOKENS HEAD_HALF idx))
(hKDisjoint : ∀ (s : BlockState)
(idx idx' : TileIndex [BLOCK_TOKENS, 1, HEAD_HALF]),
kFullFirstOffset s k_token_stride k_head_stride head_dim_stride
KV_GROUP_NUM BLOCK_TOKENS idx ≠
kFullSecondOffset s k_token_stride k_head_stride head_dim_stride
KV_GROUP_NUM BLOCK_TOKENS HEAD_HALF idx')
(v2_q_token_stride v2_q_head_stride v2_head_dim_stride v2_cos_token_stride
v2_cos_stride v2_q_total_tokens v2_Q_HEAD_NUM v2_HEAD_HALF : Nat)
(v2_k_token_stride v2_k_head_stride cacheb_stride cacheh_stride cachebs_stride
cached_stride bts_stride btb_stride block_size : Nat)
(hV2QInjFirst : ∀ s : BlockState, Function.Injective
(fun i : Fin v2_HEAD_HALF =>
v2QFirstOffset s v2_q_token_stride v2_q_head_stride v2_head_dim_stride i))
(hV2QInjSecond : ∀ s : BlockState, Function.Injective
(fun i : Fin v2_HEAD_HALF =>
v2QSecondOffset s v2_q_token_stride v2_q_head_stride v2_head_dim_stride
v2_HEAD_HALF i))
(hV2QDisjoint : ∀ (s : BlockState) (i i' : Fin v2_HEAD_HALF),
v2QFirstOffset s v2_q_token_stride v2_q_head_stride v2_head_dim_stride i ≠
v2QSecondOffset s v2_q_token_stride v2_q_head_stride v2_head_dim_stride
v2_HEAD_HALF i')
(hV2CacheDisjoint : ∀ (s : BlockState) (i i' : Fin v2_HEAD_HALF),
v2CacheOff0 s BlockTables ContextLengths cacheb_stride cacheh_stride
cachebs_stride cached_stride bts_stride btb_stride block_size
v2_q_total_tokens i ≠
v2CacheOff1 s BlockTables ContextLengths cacheb_stride cacheh_stride
cachebs_stride cached_stride bts_stride btb_stride block_size
v2_q_total_tokens v2_HEAD_HALF i') :
(∃ alg, (rotary_embedding_kernel_surface Q K Cos Sin
surf_q_token_stride surf_q_head_stride surf_k_token_stride
surf_k_head_stride surf_head_dim_stride surf_cos_token_stride
surf_cos_stride surf_q_total_tokens surf_Q_HEAD_NUM surf_KV_GROUP_NUM
surf_HEAD_DIM surf_BLOCK_TOKENS).toAlgorithm? = Except.ok alg) ∧
(rotaryNopadQIO Q Cos Sin q_token_stride q_head_stride head_dim_stride
cos_token_stride cos_stride q_total_tokens Q_HEAD_NUM HEAD_HALF BLOCK_TOKENS
⊨ fun _pid₀ _pid₁ xs o j =>
rotaryPair (xs (⟨0, by decide⟩ : Fin 4) j) (xs (⟨1, by decide⟩ : Fin 4) j)
(xs (⟨2, by decide⟩ : Fin 4) j) (xs (⟨3, by decide⟩ : Fin 4) j) o) ∧
(rotaryNopadKIO K Cos Sin k_token_stride k_head_stride head_dim_stride
cos_token_stride cos_stride k_q_total_tokens KV_GROUP_NUM HEAD_HALF BLOCK_TOKENS
⊨ fun _pid₀ _pid₁ xs o j =>
rotaryPair (xs (⟨0, by decide⟩ : Fin 4) j) (xs (⟨1, by decide⟩ : Fin 4) j)
(xs (⟨2, by decide⟩ : Fin 4) j) (xs (⟨3, by decide⟩ : Fin 4) j) o) ∧
(∃ alg, (fused_rotary_embedding_v2_surface Q K Cos Sin KVCache BlockTables
ContextLengths surf_q_token_stride surf_q_head_stride surf_k_token_stride
surf_k_head_stride surf_head_dim_stride surf_cos_token_stride surf_cos_stride
cacheb_stride cacheh_stride cachebs_stride cached_stride bts_stride btb_stride
block_size surf_q_total_tokens surf_Q_HEAD_NUM surf_HEAD_DIM).toAlgorithm?
= Except.ok alg) ∧
(fusedV2QIO Q Cos Sin v2_q_token_stride v2_q_head_stride v2_head_dim_stride
v2_cos_token_stride v2_cos_stride v2_q_total_tokens v2_Q_HEAD_NUM v2_HEAD_HALF
⊨ fun _pid₀ pid₁ xs o j =>
rotaryPair (xs (⟨0, by decide⟩ : Fin 4) j) (xs (⟨1, by decide⟩ : Fin 4) j)
(if pid₁ < v2_q_total_tokens then xs (⟨2, by decide⟩ : Fin 4) j else 0)
(if pid₁ < v2_q_total_tokens then xs (⟨3, by decide⟩ : Fin 4) j else 0) o) ∧
(ChainMetaGroupedMasked2DKernelIO.Implements
(fusedV2CacheIO KVCache K Cos Sin BlockTables ContextLengths v2_k_token_stride
v2_k_head_stride v2_head_dim_stride v2_cos_token_stride v2_cos_stride
cacheb_stride cacheh_stride cachebs_stride cached_stride bts_stride btb_stride
block_size v2_q_total_tokens v2_Q_HEAD_NUM v2_HEAD_HALF)
(fun _pid₀ pid₁ _s1 _s2 xs o j =>
rotaryPair (xs (⟨0, by decide⟩ : Fin 4) j) (xs (⟨1, by decide⟩ : Fin 4) j)
(if pid₁ < v2_q_total_tokens then xs (⟨2, by decide⟩ : Fin 4) j else 0)
(if pid₁ < v2_q_total_tokens then xs (⟨3, by decide⟩ : Fin 4) j else 0) o))
IO definition: rotaryNopadQIOdef rotaryNopadQIO (Q Cos Sin : RegionName)
(q_token_stride q_head_stride head_dim_stride cos_token_stride cos_stride
q_total_tokens Q_HEAD_NUM HEAD_HALF BLOCK_TOKENS : Nat) :
GroupedMasked2DKernelIO where
kernel := rotary_embedding_q_surface Q Cos Sin q_token_stride q_head_stride
head_dim_stride cos_token_stride cos_stride q_total_tokens Q_HEAD_NUM HEAD_HALF
BLOCK_TOKENS
nIn := 4
nOut := 2
bufs := [Q, Cos, Sin]
inp := fun i => match i with
| ⟨0, _⟩ => Q | ⟨1, _⟩ => Q | ⟨2, _⟩ => Cos | ⟨_ + 3, _⟩ => Sin
out := fun _ => Q
B := BLOCK_TOKENS * HEAD_HALF
read := fun i _pid₀ pid₁ j => match i with
| ⟨0, _⟩ => dataFirstP _pid₀ pid₁ q_token_stride q_head_stride head_dim_stride
BLOCK_TOKENS HEAD_HALF j
| ⟨1, _⟩ => dataSecondP _pid₀ pid₁ q_token_stride q_head_stride head_dim_stride
BLOCK_TOKENS HEAD_HALF j
| ⟨2, _⟩ => cosP pid₁ cos_token_stride cos_stride BLOCK_TOKENS HEAD_HALF j
| ⟨_ + 3, _⟩ => cosP pid₁ cos_token_stride cos_stride BLOCK_TOKENS HEAD_HALF j
readMask := fun i pid₀ pid₁ j => match i with
| ⟨0, _⟩ => activeQP pid₀ pid₁ q_total_tokens Q_HEAD_NUM BLOCK_TOKENS HEAD_HALF j
| ⟨1, _⟩ => activeQP pid₀ pid₁ q_total_tokens Q_HEAD_NUM BLOCK_TOKENS HEAD_HALF j
| ⟨2, _⟩ => tokP pid₁ q_total_tokens BLOCK_TOKENS HEAD_HALF j
| ⟨_ + 3, _⟩ => tokP pid₁ q_total_tokens BLOCK_TOKENS HEAD_HALF j
write := fun o pid₀ pid₁ j => match o with
| ⟨0, _⟩ => dataFirstP pid₀ pid₁ q_token_stride q_head_stride head_dim_stride
BLOCK_TOKENS HEAD_HALF j
| ⟨_ + 1, _⟩ => dataSecondP pid₀ pid₁ q_token_stride q_head_stride head_dim_stride
BLOCK_TOKENS HEAD_HALF j
writeMask := fun o pid₀ pid₁ j => match o with
| ⟨0, _⟩ => activeQP pid₀ pid₁ q_total_tokens Q_HEAD_NUM BLOCK_TOKENS HEAD_HALF j
| ⟨_ + 1, _⟩ => activeQP pid₀ pid₁ q_total_tokens Q_HEAD_NUM BLOCK_TOKENS HEAD_HALF j
IO definition: rotaryNopadKIOdef rotaryNopadKIO (K Cos Sin : RegionName)
(k_token_stride k_head_stride head_dim_stride cos_token_stride cos_stride
q_total_tokens KV_GROUP_NUM HEAD_HALF BLOCK_TOKENS : Nat) :
GroupedMasked2DKernelIO where
kernel := rotary_embedding_k_surface K Cos Sin k_token_stride k_head_stride
head_dim_stride cos_token_stride cos_stride q_total_tokens KV_GROUP_NUM HEAD_HALF
BLOCK_TOKENS
nIn := 4
nOut := 2
bufs := [K, Cos, Sin]
inp := fun i => match i with
| ⟨0, _⟩ => K | ⟨1, _⟩ => K | ⟨2, _⟩ => Cos | ⟨_ + 3, _⟩ => Sin
out := fun _ => K
B := BLOCK_TOKENS * HEAD_HALF
read := fun i pid₀ pid₁ j => match i with
| ⟨0, _⟩ => dataFirstP (pid₀ / KV_GROUP_NUM) pid₁ k_token_stride k_head_stride
head_dim_stride BLOCK_TOKENS HEAD_HALF j
| ⟨1, _⟩ => dataSecondP (pid₀ / KV_GROUP_NUM) pid₁ k_token_stride k_head_stride
head_dim_stride BLOCK_TOKENS HEAD_HALF j
| ⟨2, _⟩ => cosP pid₁ cos_token_stride cos_stride BLOCK_TOKENS HEAD_HALF j
| ⟨_ + 3, _⟩ => cosP pid₁ cos_token_stride cos_stride BLOCK_TOKENS HEAD_HALF j
readMask := fun i pid₀ pid₁ j => match i with
| ⟨0, _⟩ => activeKP pid₀ pid₁ q_total_tokens KV_GROUP_NUM BLOCK_TOKENS HEAD_HALF j
| ⟨1, _⟩ => activeKP pid₀ pid₁ q_total_tokens KV_GROUP_NUM BLOCK_TOKENS HEAD_HALF j
| ⟨2, _⟩ => activeKP pid₀ pid₁ q_total_tokens KV_GROUP_NUM BLOCK_TOKENS HEAD_HALF j
| ⟨_ + 3, _⟩ => activeKP pid₀ pid₁ q_total_tokens KV_GROUP_NUM BLOCK_TOKENS HEAD_HALF j
write := fun o pid₀ pid₁ j => match o with
| ⟨0, _⟩ => dataFirstP (pid₀ / KV_GROUP_NUM) pid₁ k_token_stride k_head_stride
head_dim_stride BLOCK_TOKENS HEAD_HALF j
| ⟨_ + 1, _⟩ => dataSecondP (pid₀ / KV_GROUP_NUM) pid₁ k_token_stride k_head_stride
head_dim_stride BLOCK_TOKENS HEAD_HALF j
writeMask := fun o pid₀ pid₁ j => match o with
| ⟨0, _⟩ => activeKP pid₀ pid₁ q_total_tokens KV_GROUP_NUM BLOCK_TOKENS HEAD_HALF j
| ⟨_ + 1, _⟩ => activeKP pid₀ pid₁ q_total_tokens KV_GROUP_NUM BLOCK_TOKENS HEAD_HALF j
IO definition: fusedV2QIOdef fusedV2QIO (Q Cos Sin : RegionName)
(q_token_stride q_head_stride head_dim_stride cos_token_stride cos_stride
q_total_tokens Q_HEAD_NUM HEAD_HALF : Nat) :
GroupedMasked2DKernelIO where
kernel := fused_rotary_embedding_v2_q_surface Q Cos Sin q_token_stride
q_head_stride head_dim_stride cos_token_stride cos_stride q_total_tokens
Q_HEAD_NUM HEAD_HALF
nIn := 4
nOut := 2
bufs := [Q, Cos, Sin]
inp := fun i => match i with
| ⟨0, _⟩ => Q | ⟨1, _⟩ => Q | ⟨2, _⟩ => Cos | ⟨_ + 3, _⟩ => Sin
out := fun _ => Q
B := HEAD_HALF
read := fun i _pid₀ pid₁ j => match i with
| ⟨0, _⟩ => pid₁ * q_token_stride + _pid₀ * q_head_stride + j.val * head_dim_stride
| ⟨1, _⟩ => pid₁ * q_token_stride + _pid₀ * q_head_stride +
(j.val + HEAD_HALF) * head_dim_stride
| ⟨2, _⟩ => pid₁ * cos_token_stride + j.val * cos_stride
| ⟨_ + 3, _⟩ => pid₁ * cos_token_stride + j.val * cos_stride
readMask := fun i pid₀ pid₁ _j => match i with
| ⟨0, _⟩ => pid₀ < Q_HEAD_NUM
| ⟨1, _⟩ => pid₀ < Q_HEAD_NUM
| ⟨2, _⟩ => pid₁ < q_total_tokens
| ⟨_ + 3, _⟩ => pid₁ < q_total_tokens
write := fun o pid₀ pid₁ j => match o with
| ⟨0, _⟩ => pid₁ * q_token_stride + pid₀ * q_head_stride + j.val * head_dim_stride
| ⟨_ + 1, _⟩ => pid₁ * q_token_stride + pid₀ * q_head_stride +
(j.val + HEAD_HALF) * head_dim_stride
writeMask := fun _o pid₀ _pid₁ _j => pid₀ < Q_HEAD_NUM
IO definition: fusedV2CacheIOdef fusedV2CacheIO (KVCache K Cos Sin : RegionName)
(BlockTables ContextLengths : Region .nat)
(k_token_stride k_head_stride head_dim_stride cos_token_stride cos_stride
cacheb_stride cacheh_stride cachebs_stride cached_stride bts_stride btb_stride
block_size q_total_tokens Q_HEAD_NUM HEAD_HALF : Nat) :
ChainMetaGroupedMasked2DKernelIO where
kernel := fused_rotary_embedding_v2_cache_surface KVCache K Cos Sin BlockTables
ContextLengths k_token_stride k_head_stride head_dim_stride cos_token_stride
cos_stride cacheb_stride cacheh_stride cachebs_stride cached_stride bts_stride
btb_stride block_size q_total_tokens Q_HEAD_NUM HEAD_HALF
nIn := 4
nOut := 2
bufs := [KVCache, K, Cos, Sin, BlockTables, ContextLengths]
mbuf1 := ContextLengths
mbuf2 := BlockTables
inp := fun i => match i with
| ⟨0, _⟩ => K | ⟨1, _⟩ => K | ⟨2, _⟩ => Cos | ⟨_ + 3, _⟩ => Sin
out := fun _ => KVCache
B := HEAD_HALF
mwin1 := fun _pid₀ pid₁ => pid₁
mwin2 := fun _pid₀ pid₁ s1 =>
pid₁ * bts_stride + (s1 - 1) / block_size * btb_stride
read := fun i pid₀ pid₁ _s1 _s2 j => match i with
| ⟨0, _⟩ => pid₁ * k_token_stride + pid₀ * k_head_stride + j.val * head_dim_stride
| ⟨1, _⟩ => pid₁ * k_token_stride + pid₀ * k_head_stride +
(j.val + HEAD_HALF) * head_dim_stride
| ⟨2, _⟩ => pid₁ * cos_token_stride + j.val * cos_stride
| ⟨_ + 3, _⟩ => pid₁ * cos_token_stride + j.val * cos_stride
readMask := fun i pid₀ pid₁ _s1 _s2 _j => match i with
| ⟨0, _⟩ => pid₀ < Q_HEAD_NUM
| ⟨1, _⟩ => pid₀ < Q_HEAD_NUM
| ⟨2, _⟩ => pid₁ < q_total_tokens
| ⟨_ + 3, _⟩ => pid₁ < q_total_tokens
write := fun o pid₀ pid₁ s1 s2 j => match o with
| ⟨0, _⟩ => (if pid₁ < q_total_tokens then s2 * cacheb_stride else 0) +
pid₀ * cacheh_stride + (s1 - 1) % block_size * cachebs_stride +
j.val * cached_stride
| ⟨_ + 1, _⟩ => (if pid₁ < q_total_tokens then s2 * cacheb_stride else 0) +
pid₀ * cacheh_stride + (s1 - 1) % block_size * cachebs_stride +
(j.val + HEAD_HALF) * cached_stride
writeMask := fun _o pid₀ _pid₁ _s1 _s2 _j => pid₀ < Q_HEAD_NUM
Referenced model/specification definitions: Python source |
|---|
rotary_transformrotary_transform_python_output_summary_generalMathematical execution | Configured model / stage Selected row/body rotary value facts and a metadata-aware row IO contract. Interleaving, varlen, offsets, and paired-half layout conditions must be read from the chosen statement. | Statement and preconditionsspecification rotary_transform_python_output_summary_general
(OUT X COS SIN : RegionName) (CU_SEQLENS SEQLEN_OFFSETS_REGION : Region .nat)
(SEQLEN_OFFSETS_SCALAR seqlen_surf nheads rotary_dim seqlen_ro_surf CACHE_KEY_SEQLEN
surf_stride_out_batch surf_stride_out_seqlen surf_stride_out_nheads surf_stride_out_headdim
surf_stride_x_batch surf_stride_x_seqlen surf_stride_x_nheads surf_stride_x_headdim
BLOCK_K BLOCK_M_surf : Nat)
(IS_SEQLEN_OFFSETS_TENSOR IS_VARLEN INTERLEAVED CONJUGATE : Bool)
(SEQLEN_OFFSETS seqlen rotary_dim_half seqlen_ro
stride_out_batch stride_out_seqlen stride_out_nheads stride_out_headdim
stride_x_batch stride_x_seqlen stride_x_nheads stride_x_headdim
BLOCK_M BLOCK_HALF : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_HALF =>
outOffset s stride_out_batch stride_out_seqlen stride_out_nheads
stride_out_headdim BLOCK_M i))
(hOut1Inj : Function.Injective
(fun i : Fin BLOCK_HALF =>
out1Offset s stride_out_batch stride_out_seqlen stride_out_nheads
stride_out_headdim rotary_dim_half BLOCK_M i))
(hStrideHd : stride_out_headdim ≠ 0)
(hHalfBound : BLOCK_HALF ≤ rotary_dim_half) :
(∃ alg, (rotary_kernel_surface OUT X COS SIN CU_SEQLENS SEQLEN_OFFSETS_REGION
SEQLEN_OFFSETS_SCALAR seqlen_surf nheads rotary_dim seqlen_ro_surf CACHE_KEY_SEQLEN
surf_stride_out_batch surf_stride_out_seqlen surf_stride_out_nheads surf_stride_out_headdim
surf_stride_x_batch surf_stride_x_seqlen surf_stride_x_nheads surf_stride_x_headdim
BLOCK_K BLOCK_M_surf IS_SEQLEN_OFFSETS_TENSOR IS_VARLEN INTERLEAVED
CONJUGATE).toAlgorithm? = Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := rotary_kernel_o0o1_row OUT X COS SIN SEQLEN_OFFSETS
seqlen rotary_dim_half seqlen_ro stride_out_batch stride_out_seqlen
stride_out_nheads stride_out_headdim stride_x_batch stride_x_seqlen
stride_x_nheads stride_x_headdim BLOCK_M BLOCK_HALF)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_HALF => active s seqlen rotary_dim_half BLOCK_M i)
(fun i => (OUT,
outOffset s stride_out_batch stride_out_seqlen stride_out_nheads
stride_out_headdim BLOCK_M i)))
(expected := fun i =>
rotaryO0Spec s X COS SIN SEQLEN_OFFSETS seqlen_ro stride_x_batch
stride_x_seqlen stride_x_nheads stride_x_headdim rotary_dim_half
BLOCK_M i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := rotary_kernel_o0o1_row OUT X COS SIN SEQLEN_OFFSETS
seqlen rotary_dim_half seqlen_ro stride_out_batch stride_out_seqlen
stride_out_nheads stride_out_headdim stride_x_batch stride_x_seqlen
stride_x_nheads stride_x_headdim BLOCK_M BLOCK_HALF)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_HALF => active s seqlen rotary_dim_half BLOCK_M i)
(fun i => (OUT,
out1Offset s stride_out_batch stride_out_seqlen stride_out_nheads
stride_out_headdim rotary_dim_half BLOCK_M i)))
(expected := fun i =>
rotaryO1Spec s X COS SIN SEQLEN_OFFSETS seqlen_ro stride_x_batch
stride_x_seqlen stride_x_nheads stride_x_headdim rotary_dim_half
BLOCK_M i))
Referenced model/specification definitions: Python source |
|---|
rotary_transformrotary_transform_meta_implementsMathematical execution | Configured model / stage Selected row/body rotary value facts and a metadata-aware row IO contract. Interleaving, varlen, offsets, and paired-half layout conditions must be read from the chosen statement. | Statement and preconditionsspecification rotary_transform_meta_implements
(OUT X COS SIN : RegionName) (CU_SEQLENS : Region .nat)
(HEAD_IDX SEQLEN_OFFSETS seqlen rotary_dim_half seqlen_ro
stride_out_batch stride_out_seqlen stride_out_nheads stride_out_headdim
stride_x_batch stride_x_seqlen stride_x_nheads stride_x_headdim
BLOCK_M BLOCK_HALF : Nat) (IS_VARLEN : Bool)
(hStrideHd : stride_out_headdim ≠ 0)
(hHalfBound : BLOCK_HALF ≤ rotary_dim_half) :
rotaryMetaIO OUT X COS SIN CU_SEQLENS HEAD_IDX SEQLEN_OFFSETS seqlen rotary_dim_half
seqlen_ro stride_out_batch stride_out_seqlen stride_out_nheads stride_out_headdim
stride_x_batch stride_x_seqlen stride_x_nheads stride_x_headdim BLOCK_M BLOCK_HALF
IS_VARLEN
⊨ fun pid₀ _pid₁ _s1 _s2 xs o j =>
rotaryMetaF SEQLEN_OFFSETS seqlen_ro rotary_dim_half BLOCK_M pid₀ xs o j
IO definition: rotaryMetaIOdef rotaryMetaIO
(OUT X COS SIN : RegionName) (CU_SEQLENS : Region .nat)
(HEAD_IDX SEQLEN_OFFSETS seqlen rotary_dim_half seqlen_ro
stride_out_batch stride_out_seqlen stride_out_nheads stride_out_headdim
stride_x_batch stride_x_seqlen stride_x_nheads stride_x_headdim
BLOCK_M BLOCK_HALF : Nat) (IS_VARLEN : Bool) :
MetaGroupedMasked2DKernelIO where
kernel := rotary_meta_row OUT X COS SIN CU_SEQLENS HEAD_IDX SEQLEN_OFFSETS seqlen
rotary_dim_half seqlen_ro stride_out_batch stride_out_seqlen stride_out_nheads
stride_out_headdim stride_x_batch stride_x_seqlen stride_x_nheads stride_x_headdim
BLOCK_M BLOCK_HALF IS_VARLEN
nIn := 4
nOut := 2
bufs := [X, OUT, COS, SIN, CU_SEQLENS]
mbuf1 := CU_SEQLENS
mbuf2 := CU_SEQLENS
inp := fun i => match i with
| ⟨0, _⟩ => X
| ⟨1, _⟩ => X
| ⟨2, _⟩ => COS
| ⟨_ + 3, _⟩ => SIN
out := fun _ => OUT
B := BLOCK_HALF
mwin1 := fun _ pid₁ => pid₁
mwin2 := fun _ pid₁ => pid₁ + 1
read := fun i pid₀ pid₁ s1 _s2 j => match i with
| ⟨0, _⟩ => ioXBase IS_VARLEN pid₁ s1 HEAD_IDX stride_x_batch stride_x_seqlen
stride_x_nheads + pid₀ * BLOCK_M * stride_x_seqlen + j.val * stride_x_headdim
| ⟨1, _⟩ => ioXBase IS_VARLEN pid₁ s1 HEAD_IDX stride_x_batch stride_x_seqlen
stride_x_nheads + pid₀ * BLOCK_M * stride_x_seqlen +
(j.val + rotary_dim_half) * stride_x_headdim
| ⟨2, _⟩ => (pid₀ * BLOCK_M + SEQLEN_OFFSETS) * rotary_dim_half + j.val
| ⟨_ + 3, _⟩ => (pid₀ * BLOCK_M + SEQLEN_OFFSETS) * rotary_dim_half + j.val
readMask := fun i pid₀ _pid₁ s1 s2 j => match i with
| ⟨0, _⟩ => pid₀ * BLOCK_M < ioSeqlen IS_VARLEN s1 s2 seqlen ∧ j.val < rotary_dim_half
| ⟨1, _⟩ => pid₀ * BLOCK_M < ioSeqlen IS_VARLEN s1 s2 seqlen ∧ j.val < rotary_dim_half
| ⟨2, _⟩ => pid₀ * BLOCK_M + SEQLEN_OFFSETS < seqlen_ro ∧ j.val < rotary_dim_half
| ⟨_ + 3, _⟩ => pid₀ * BLOCK_M + SEQLEN_OFFSETS < seqlen_ro ∧ j.val < rotary_dim_half
write := fun o pid₀ pid₁ s1 _s2 j =>
ioOutBase IS_VARLEN pid₁ s1 HEAD_IDX stride_out_batch stride_out_seqlen
stride_out_nheads + pid₀ * BLOCK_M * stride_out_seqlen +
(j.val + (match o with | ⟨0, _⟩ => 0 | ⟨_ + 1, _⟩ => rotary_dim_half)) *
stride_out_headdim
writeMask := fun _ pid₀ _pid₁ s1 s2 j =>
pid₀ * BLOCK_M < ioSeqlen IS_VARLEN s1 s2 seqlen ∧ j.val < rotary_dim_half
Referenced model/specification definitions: Python source |
|---|
rotary_transform_opsrotary_transform_ops_output_summary_generalMathematical execution | Configured model / stage Selected rotary body/row facts and metadata-aware IO contract. Flags, paired-half addresses, and variable-length metadata restrict the configured Python path. | Statement and preconditionsspecification rotary_transform_ops_output_summary_general
(OUT X COS SIN : RegionName) (CU_SEQLENS SEQLEN_OFFSETS : Region .nat)
(SEQLEN_OFFSETS_SCALAR seqlen rotary_dim seqlen_ro
stride_out_batch stride_out_seqlen stride_out_nheads stride_out_headdim
stride_x_batch stride_x_seqlen stride_x_nheads stride_x_headdim
BLOCK_K BLOCK_M : Nat)
(IS_SEQLEN_OFFSETS_TENSOR IS_VARLEN INTERLEAVED CONJUGATE : Bool)
(body_SEQLEN_OFFSETS body_seqlen body_rotary_dim_half body_seqlen_ro
body_stride_out_batch body_stride_out_seqlen body_stride_out_nheads
body_stride_out_headdim body_stride_x_batch body_stride_x_seqlen
body_stride_x_nheads body_stride_x_headdim body_BLOCK_M BLOCK_HALF : Nat)
(s : BlockState)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_HALF =>
outOffset s body_stride_out_batch body_stride_out_seqlen
body_stride_out_nheads body_stride_out_headdim body_BLOCK_M i))
(hOut1Inj : Function.Injective
(fun i : Fin BLOCK_HALF =>
out1Offset s body_stride_out_batch body_stride_out_seqlen
body_stride_out_nheads body_stride_out_headdim body_rotary_dim_half
body_BLOCK_M i))
(hStrideHd : body_stride_out_headdim ≠ 0)
(hHalfBound : BLOCK_HALF ≤ body_rotary_dim_half) :
(∃ alg, (rotary_kernel_surface OUT X COS SIN CU_SEQLENS SEQLEN_OFFSETS
SEQLEN_OFFSETS_SCALAR seqlen rotary_dim seqlen_ro stride_out_batch
stride_out_seqlen stride_out_nheads stride_out_headdim stride_x_batch
stride_x_seqlen stride_x_nheads stride_x_headdim BLOCK_K BLOCK_M
IS_SEQLEN_OFFSETS_TENSOR IS_VARLEN INTERLEAVED CONJUGATE).toAlgorithm? =
Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := rotary_kernel_o0o1_row OUT X COS SIN
body_SEQLEN_OFFSETS body_seqlen body_rotary_dim_half body_seqlen_ro
body_stride_out_batch body_stride_out_seqlen body_stride_out_nheads
body_stride_out_headdim body_stride_x_batch body_stride_x_seqlen
body_stride_x_nheads body_stride_x_headdim body_BLOCK_M BLOCK_HALF)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_HALF =>
active s body_seqlen body_rotary_dim_half body_BLOCK_M i)
(fun i => (OUT,
outOffset s body_stride_out_batch body_stride_out_seqlen
body_stride_out_nheads body_stride_out_headdim body_BLOCK_M i)))
(expected := fun i =>
rotaryO0Spec s X COS SIN body_SEQLEN_OFFSETS body_seqlen_ro
body_stride_x_batch body_stride_x_seqlen body_stride_x_nheads
body_stride_x_headdim body_rotary_dim_half body_BLOCK_M i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := rotary_kernel_o0o1_row OUT X COS SIN
body_SEQLEN_OFFSETS body_seqlen body_rotary_dim_half body_seqlen_ro
body_stride_out_batch body_stride_out_seqlen body_stride_out_nheads
body_stride_out_headdim body_stride_x_batch body_stride_x_seqlen
body_stride_x_nheads body_stride_x_headdim body_BLOCK_M BLOCK_HALF)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_HALF =>
active s body_seqlen body_rotary_dim_half body_BLOCK_M i)
(fun i => (OUT,
out1Offset s body_stride_out_batch body_stride_out_seqlen
body_stride_out_nheads body_stride_out_headdim body_rotary_dim_half
body_BLOCK_M i)))
(expected := fun i =>
rotaryO1Spec s X COS SIN body_SEQLEN_OFFSETS body_seqlen_ro
body_stride_x_batch body_stride_x_seqlen body_stride_x_nheads
body_stride_x_headdim body_rotary_dim_half body_BLOCK_M i))
Referenced model/specification definitions: Python source |
|---|
rotary_transform_opsrotary_transform_ops_meta_implementsMathematical execution | Configured model / stage Selected rotary body/row facts and metadata-aware IO contract. Flags, paired-half addresses, and variable-length metadata restrict the configured Python path. | Statement and preconditionsspecification rotary_transform_ops_meta_implements
(OUT X COS SIN : RegionName) (CU_SEQLENS : Region .nat)
(HEAD_IDX SEQLEN_OFFSETS seqlen rotary_dim_half seqlen_ro
stride_out_batch stride_out_seqlen stride_out_nheads stride_out_headdim
stride_x_batch stride_x_seqlen stride_x_nheads stride_x_headdim
BLOCK_M BLOCK_HALF : Nat) (IS_VARLEN : Bool)
(hStrideHd : stride_out_headdim ≠ 0)
(hHalfBound : BLOCK_HALF ≤ rotary_dim_half) :
rotaryMetaIO OUT X COS SIN CU_SEQLENS HEAD_IDX SEQLEN_OFFSETS seqlen rotary_dim_half
seqlen_ro stride_out_batch stride_out_seqlen stride_out_nheads stride_out_headdim
stride_x_batch stride_x_seqlen stride_x_nheads stride_x_headdim BLOCK_M BLOCK_HALF
IS_VARLEN
⊨ fun pid₀ _pid₁ _s1 _s2 xs o j =>
rotaryMetaF SEQLEN_OFFSETS seqlen_ro rotary_dim_half BLOCK_M pid₀ xs o j
IO definition: rotaryMetaIOdef rotaryMetaIO
(OUT X COS SIN : RegionName) (CU_SEQLENS : Region .nat)
(HEAD_IDX SEQLEN_OFFSETS seqlen rotary_dim_half seqlen_ro
stride_out_batch stride_out_seqlen stride_out_nheads stride_out_headdim
stride_x_batch stride_x_seqlen stride_x_nheads stride_x_headdim
BLOCK_M BLOCK_HALF : Nat) (IS_VARLEN : Bool) :
MetaGroupedMasked2DKernelIO where
kernel := rotary_meta_row OUT X COS SIN CU_SEQLENS HEAD_IDX SEQLEN_OFFSETS seqlen
rotary_dim_half seqlen_ro stride_out_batch stride_out_seqlen stride_out_nheads
stride_out_headdim stride_x_batch stride_x_seqlen stride_x_nheads stride_x_headdim
BLOCK_M BLOCK_HALF IS_VARLEN
nIn := 4
nOut := 2
bufs := [X, OUT, COS, SIN, CU_SEQLENS]
mbuf1 := CU_SEQLENS
mbuf2 := CU_SEQLENS
inp := fun i => match i with
| ⟨0, _⟩ => X
| ⟨1, _⟩ => X
| ⟨2, _⟩ => COS
| ⟨_ + 3, _⟩ => SIN
out := fun _ => OUT
B := BLOCK_HALF
mwin1 := fun _ pid₁ => pid₁
mwin2 := fun _ pid₁ => pid₁ + 1
read := fun i pid₀ pid₁ s1 _s2 j => match i with
| ⟨0, _⟩ => ioXBase IS_VARLEN pid₁ s1 HEAD_IDX stride_x_batch stride_x_seqlen
stride_x_nheads + pid₀ * BLOCK_M * stride_x_seqlen + j.val * stride_x_headdim
| ⟨1, _⟩ => ioXBase IS_VARLEN pid₁ s1 HEAD_IDX stride_x_batch stride_x_seqlen
stride_x_nheads + pid₀ * BLOCK_M * stride_x_seqlen +
(j.val + rotary_dim_half) * stride_x_headdim
| ⟨2, _⟩ => (pid₀ * BLOCK_M + SEQLEN_OFFSETS) * rotary_dim_half + j.val
| ⟨_ + 3, _⟩ => (pid₀ * BLOCK_M + SEQLEN_OFFSETS) * rotary_dim_half + j.val
readMask := fun i pid₀ _pid₁ s1 s2 j => match i with
| ⟨0, _⟩ => pid₀ * BLOCK_M < ioSeqlen IS_VARLEN s1 s2 seqlen ∧ j.val < rotary_dim_half
| ⟨1, _⟩ => pid₀ * BLOCK_M < ioSeqlen IS_VARLEN s1 s2 seqlen ∧ j.val < rotary_dim_half
| ⟨2, _⟩ => pid₀ * BLOCK_M + SEQLEN_OFFSETS < seqlen_ro ∧ j.val < rotary_dim_half
| ⟨_ + 3, _⟩ => pid₀ * BLOCK_M + SEQLEN_OFFSETS < seqlen_ro ∧ j.val < rotary_dim_half
write := fun o pid₀ pid₁ s1 _s2 j =>
ioOutBase IS_VARLEN pid₁ s1 HEAD_IDX stride_out_batch stride_out_seqlen
stride_out_nheads + pid₀ * BLOCK_M * stride_out_seqlen +
(j.val + (match o with | ⟨0, _⟩ => 0 | ⟨_ + 1, _⟩ => rotary_dim_half)) *
stride_out_headdim
writeMask := fun _ pid₀ _pid₁ s1 s2 j =>
pid₀ * BLOCK_M < ioSeqlen IS_VARLEN s1 s2 seqlen ∧ j.val < rotary_dim_half
Referenced model/specification definitions: Python source |
|---|
rowwise_quantization_tritonquantize_rowwise_correctnessMathematical execution | Pre-rounding slice Computes row maximum and real-valued scaled outputs, but omits the faithful CUDA llrint/final-int8 result. The max producer is included; final rounding is not. | Statement and preconditionsspecification quantize_rowwise_correctness
(x_ptr inert output_ptr output_maxs : RegionName)
(n_elements BLOCK_SIZE P2 : Nat) (hP : 0 < P2)
(hRegions : output_ptr ≠ output_maxs) :
quantizeRowwiseIO x_ptr inert output_ptr output_maxs n_elements
BLOCK_SIZE P2 ⊨
fun _ _ xs _ =>
(fun i => qrOutSpec BLOCK_SIZE P2 xs i,
fun _ => qrMaxSpec BLOCK_SIZE P2 xs)
IO definition: quantizeRowwiseIOdef quantizeRowwiseIO (x_ptr inert output_ptr output_maxs : RegionName)
(n_elements BLOCK_SIZE P2 : Nat) : Masked2DKernelIO₂ₓ₂ where
kernel := quantize_rowwise_scaled x_ptr output_ptr output_maxs n_elements
BLOCK_SIZE P2
in1 := x_ptr
in2 := inert
out1 := output_ptr
out2 := output_maxs
B := P2
read1 := fun pid _ j => pid * BLOCK_SIZE + j.val
read2 := fun _ _ _ => 0
write1 := fun pid _ j => pid * BLOCK_SIZE + j.val
write2 := fun pid _ _ => pid
mask := fun _ _ j => j.val < BLOCK_SIZE
read2Mask := fun _ _ _ => False
writeMask2 := fun _ _ j => j.val = 0
Referenced model/specification definitions: Python source |
|---|
sgmv_expand_slicesgmv_expand_slice_one_row_block_output_summaryMathematical execution | Configured model / stage One sequence-row block of LoRA slice expansion with K=BLOCK_K*numKBlocks, supplied sequence/adapter metadata, and explicit output-layout premises. | Statement and preconditionsspecification sgmv_expand_slice_one_row_block_output_summary
(input_ptr lora_ptr out_ptr : RegionName)
(b_seq_start_loc seq_lens lora_indices : Region .nat) (s : BlockState)
(N xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride slice_offset BLOCK_M BLOCK_N BLOCK_K numKBlocks : Nat) (hBK : 0 < BLOCK_K)
(hInj : Function.Injective (cOffset s b_seq_start_loc cm_stride cn_stride slice_offset BLOCK_M BLOCK_N))
(hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (sgmv_expand_slice_surface input_ptr lora_ptr out_ptr b_seq_start_loc seq_lens
lora_indices N (BLOCK_K * numKBlocks) xm_stride xk_stride l0_stride lora_k_stride
lora_n_stride cm_stride cn_stride slice_offset BLOCK_M BLOCK_N BLOCK_K).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := sgmv_expand_slice_surface input_ptr lora_ptr out_ptr b_seq_start_loc seq_lens
lora_indices N (BLOCK_K * numKBlocks) xm_stride xk_stride l0_stride lora_k_stride
lora_n_stride cm_stride cn_stride slice_offset BLOCK_M BLOCK_N BLOCK_K)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_N] => activeLane s seq_lens N BLOCK_M BLOCK_N idx)
(fun idx => (out_ptr, cOffset s b_seq_start_loc cm_stride cn_stride slice_offset BLOCK_M BLOCK_N idx)))
(expected := fun idx =>
sgmvSpec s input_ptr lora_ptr b_seq_start_loc seq_lens lora_indices
N xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
BLOCK_M BLOCK_N BLOCK_K numKBlocks idx.1 idx.2.1)
Referenced model/specification definitions: Python source |
|---|
sgmv_expand_slicesgmv_expand_slice_io_correctnessAbstract cast/store rounding | Configured model / stage One sequence-row block of LoRA slice expansion with K=BLOCK_K*numKBlocks, supplied sequence/adapter metadata, and explicit output-layout premises. | Statement and preconditionsspecification sgmv_expand_slice_io_correctness (R : RoundingModel)
(input_ptr lora_ptr out_ptr : RegionName)
(b_seq_start_loc seq_lens lora_indices : Region .nat)
(N xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride slice_offset BLOCK_M BLOCK_N BLOCK_K numKBlocks : Nat)
(hBK : 0 < BLOCK_K)
(hInj : ∀ pid₀ pid₁ start : Nat,
Function.Injective (fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
(start + (pid₀ * BLOCK_M + idx.1.val)) * cm_stride
+ ((pid₁ * BLOCK_N + idx.2.1.val) + slice_offset) * cn_stride)) :
sgmvExpandSliceIO input_ptr lora_ptr out_ptr b_seq_start_loc seq_lens lora_indices
N xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride slice_offset BLOCK_M BLOCK_N BLOCK_K numKBlocks ⊨[R]
fun _ _ _ _ xs ys l =>
∑ t : Fin numKBlocks, ∑ e : Fin BLOCK_K,
xs t (aLane BLOCK_M BLOCK_N BLOCK_K l e)
* ys t (bLane BLOCK_M BLOCK_N BLOCK_K l e)
IO definition: sgmvExpandSliceIOdef sgmvExpandSliceIO (input_ptr lora_ptr out_ptr : RegionName)
(b_seq_start_loc seq_lens lora_indices : Region .nat)
(N xm_stride xk_stride l0_stride lora_k_stride lora_n_stride
cm_stride cn_stride slice_offset BLOCK_M BLOCK_N BLOCK_K numKBlocks : Nat) :
StreamMetaMasked3DKernelIO₂ where
kernel := sgmv_expand_slice_surface input_ptr lora_ptr out_ptr b_seq_start_loc seq_lens
lora_indices N (BLOCK_K * numKBlocks) xm_stride xk_stride l0_stride lora_k_stride
lora_n_stride cm_stride cn_stride slice_offset BLOCK_M BLOCK_N BLOCK_K
inp1 := input_ptr
inp2 := lora_ptr
out := out_ptr
nMeta := 3
sty := fun _ => ChanTy.nat
mbuf := sgmvMetaBuf b_seq_start_loc seq_lens lora_indices
mwin := fun _ _ _ pid₂ => pid₂
T := numKBlocks
B1 := BLOCK_M * BLOCK_K
B2 := BLOCK_K * BLOCK_N
C := BLOCK_M * BLOCK_N
read1 := fun pid₀ _ _ m t j =>
m (⟨1, by omega⟩ : Fin 3) * xm_stride
+ ((pid₀ * BLOCK_M + j.val / BLOCK_K) % m (⟨0, by omega⟩ : Fin 3)) * xm_stride
+ (t.val * BLOCK_K + j.val % BLOCK_K) * xk_stride
read2 := fun _ pid₁ _ m t j =>
l0_stride * m (⟨2, by omega⟩ : Fin 3)
+ (t.val * BLOCK_K + j.val / BLOCK_N) * lora_n_stride
+ ((pid₁ * BLOCK_N + j.val % BLOCK_N) % N) * lora_k_stride
write := fun pid₀ pid₁ _ m j =>
(m (⟨1, by omega⟩ : Fin 3) + (pid₀ * BLOCK_M + j.val / BLOCK_N)) * cm_stride
+ ((pid₁ * BLOCK_N + j.val % BLOCK_N) + slice_offset) * cn_stride
mask1 := fun _ _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ _ => True
writeMask := fun pid₀ pid₁ _ m j =>
pid₀ * BLOCK_M + j.val / BLOCK_N < m (⟨0, by omega⟩ : Fin 3)
∧ pid₁ * BLOCK_N + j.val % BLOCK_N < N
Referenced model/specification definitions: Python source |
|---|
sin_computationsin_kernel_correctnessMathematical execution | Configured model / stage Masked per-program mathematical sine. Host launch and the hardware sine approximation are outside the contract. | Statement and preconditionsspecification sin_kernel_correctness
(in_ptr0 out_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) :
sinIO in_ptr0 out_ptr n_elements BLOCK_SIZE
⊨ fun xs i => Real.sin (xs i)
IO definition: sinIOdef sinIO (in_ptr0 out_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) : MaskedKernelIO₁ where
kernel := sin_kernel in_ptr0 out_ptr n_elements BLOCK_SIZE
inp := in_ptr0
out := out_ptr
B := BLOCK_SIZE
read := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < n_elements
Referenced model/specification definitions: Python source |
|---|
sin_kernelkernel_function_correctnessMathematical execution | Configured model / stage Masked per-program mathematical sine. The theorem does not establish a hardware transcendental error bound or host launch correctness. | Statement and preconditionsspecification kernel_function_correctness
(x_ptr output_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) :
sinIO x_ptr output_ptr n_elements BLOCK_SIZE
⊨ fun xs i => Real.sin (xs i)
IO definition: sinIOdef sinIO (x_ptr output_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) : MaskedKernelIO₁ where
kernel := kernel_function x_ptr output_ptr n_elements BLOCK_SIZE
inp := x_ptr
out := output_ptr
B := BLOCK_SIZE
read := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < n_elements
Referenced model/specification definitions: Python source |
|---|
softmax_flaggemssoftmax_kernel_inner_one_tile_compute_correctMathematical execution | Configured model / stage One-tile inner/non-inner forward and backward branch specializations. The IO theorem is only the inner forward one-tile kernel; multi-tile online modes are not implied. | Statement and preconditionsspecification softmax_kernel_inner_one_tile_compute_correct
(output_ptr input_ptr : RegionName)
(N TILE_N : Nat)
(s : BlockState) :
ComputeCorrect.Realizes_without_Rounding
(kernel := softmax_kernel_inner_one_tile output_ptr input_ptr N TILE_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin TILE_N => i.val < N)
(fun i => (output_ptr, linearOffset s N i)))
(expected := fun i => softmaxFlaggemsSpec s input_ptr N TILE_N i)
Referenced model/specification definitions: Python source |
|---|
softmax_flaggemssoftmax_kernel_non_inner_one_tile_compute_correctMathematical execution | Configured model / stage One-tile inner/non-inner forward and backward branch specializations. The IO theorem is only the inner forward one-tile kernel; multi-tile online modes are not implied. | Statement and preconditionsspecification softmax_kernel_non_inner_one_tile_compute_correct
(output_ptr input_ptr : RegionName)
(N K TILE_N TILE_K : Nat)
(s : BlockState)
(hRange : s.pids 1 * TILE_K + TILE_K ≤ K) :
ComputeCorrect.Realizes_without_Rounding
(kernel := softmax_kernel_non_inner_one_tile_surface
output_ptr input_ptr N K TILE_N TILE_K)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [TILE_N, TILE_K] =>
idx.1.val < N ∧ s.pids 1 * TILE_K + idx.2.1.val < K)
(fun idx => (output_ptr, nonInnerOffset s N K TILE_K idx)))
(expected := fun idx =>
softmaxFlaggemsNonInnerSpec s input_ptr N K TILE_N TILE_K idx)
Referenced model/specification definitions: Python source |
|---|
softmax_flaggemssoftmax_backward_kernel_inner_one_tile_compute_correctMathematical execution | Configured model / stage One-tile inner/non-inner forward and backward branch specializations. The IO theorem is only the inner forward one-tile kernel; multi-tile online modes are not implied. | Statement and preconditionsspecification softmax_backward_kernel_inner_one_tile_compute_correct
(out_ptr out_grad_ptr in_grad_ptr : RegionName)
(M N TILE_M TILE_N : Nat)
(s : BlockState)
(hOffInj : Function.Injective
(fun idx : TileIndex [TILE_M, TILE_N] => innerBwdOffset s N TILE_M idx)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := softmax_backward_kernel_inner_one_tile_surface
out_ptr out_grad_ptr in_grad_ptr M N TILE_M TILE_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [TILE_M, TILE_N] =>
innerBwdActive s M N TILE_M idx)
(fun idx => (in_grad_ptr, innerBwdOffset s N TILE_M idx)))
(expected := fun idx =>
innerBwdSpec s out_ptr out_grad_ptr M N TILE_M TILE_N idx)
Referenced model/specification definitions: Python source |
|---|
softmax_flaggemssoftmax_backward_kernel_non_inner_one_tile_compute_correctMathematical execution | Configured model / stage One-tile inner/non-inner forward and backward branch specializations. The IO theorem is only the inner forward one-tile kernel; multi-tile online modes are not implied. | Statement and preconditionsspecification softmax_backward_kernel_non_inner_one_tile_compute_correct
(out_ptr out_grad_ptr in_grad_ptr : RegionName)
(N K TILE_N TILE_K : Nat)
(s : BlockState)
(hRange : s.pids 1 * TILE_K + TILE_K ≤ K) :
ComputeCorrect.Realizes_without_Rounding
(kernel := softmax_backward_kernel_non_inner_one_tile_surface
out_ptr out_grad_ptr in_grad_ptr N K TILE_N TILE_K)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [TILE_N, TILE_K] =>
idx.1.val < N ∧ s.pids 1 * TILE_K + idx.2.1.val < K)
(fun idx => (in_grad_ptr, nonInnerOffset s N K TILE_K idx)))
(expected := fun idx =>
nonInnerBwdSpec s out_ptr out_grad_ptr N K TILE_N TILE_K idx)
Referenced model/specification definitions: Python source |
|---|
softmax_flaggemssoftmax_kernel_inner_one_tile_correctnessMathematical execution | Configured model / stage One-tile inner/non-inner forward and backward branch specializations. The IO theorem is only the inner forward one-tile kernel; multi-tile online modes are not implied. | Statement and preconditionsspecification softmax_kernel_inner_one_tile_correctness
(output_ptr input_ptr : RegionName) (N TILE_N : Nat) (hT : 0 < TILE_N) :
softmaxFlaggemsInnerIO output_ptr input_ptr N TILE_N
⊨ fun xs i => softmaxFlaggemsSpecOf N TILE_N xs i
IO definition: softmaxFlaggemsInnerIOdef softmaxFlaggemsInnerIO (output_ptr input_ptr : RegionName)
(N TILE_N : Nat) : MaskedKernelIO₁ where
kernel := softmax_kernel_inner_one_tile output_ptr input_ptr N TILE_N
inp := input_ptr
out := output_ptr
B := TILE_N
read := fun pid => pid * N
write := fun pid => pid * N
mask := fun _ j => j.val < N
Referenced model/specification definitions: Python source |
|---|
softmax_optimizesoftmax_kernel_online_v2_one_tile_output_summaryMathematical execution | Configured model / stage Separate one-tile and full online-v2 value targets. The IO theorem covers only the one-tile surface; inspect block/row assumptions before using the full value theorem. | Statement and preconditionsspecification softmax_kernel_online_v2_one_tile_output_summary
(output_ptr input_ptr : RegionName)
(N TILE_N : Nat)
(s : BlockState) :
(∃ alg, (softmax_kernel_online_v2_one_tile output_ptr input_ptr N TILE_N).toAlgorithm? =
Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := softmax_kernel_online_v2_one_tile output_ptr input_ptr N TILE_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin TILE_N => i.val < N)
(fun i => (output_ptr, linearOffset s N i)))
(expected := fun i => softmaxOptimizeSpec s input_ptr N TILE_N i)
Referenced model/specification definitions: Python source |
|---|
softmax_optimizesoftmax_kernel_online_v2_output_summaryMathematical execution | Configured model / stage Separate one-tile and full online-v2 value targets. The IO theorem covers only the one-tile surface; inspect block/row assumptions before using the full value theorem. | Statement and preconditionsspecification softmax_kernel_online_v2_output_summary
(output_ptr input_ptr : RegionName) (M N TILE_N : Nat)
(hN : 0 < N) (hT : 0 < TILE_N) (hne : output_ptr ≠ input_ptr)
(s : BlockState) :
ComputeCorrect.Realizes_without_Rounding
(kernel := softmax_kernel_online_v2_surface output_ptr input_ptr M N TILE_N)
(initialState := s)
(write := fun j : Fin N => some (output_ptr, linearOffset s N j))
(expected := fun j : Fin N => softmaxOptimizeFullSpec s input_ptr N j)
Referenced model/specification definitions: Python source |
|---|
softmax_optimizesoftmax_kernel_online_v2_one_tile_io_correctnessMathematical execution | Configured model / stage Separate one-tile and full online-v2 value targets. The IO theorem covers only the one-tile surface; inspect block/row assumptions before using the full value theorem. | Statement and preconditionsspecification softmax_kernel_online_v2_one_tile_io_correctness
(output_ptr input_ptr : RegionName) (N TILE_N : Nat) (hT : 0 < TILE_N) :
oneTileIO output_ptr input_ptr N TILE_N
⊨ fun _pid xs idx => softmaxTileSpec N TILE_N xs idx.1
IO definition: oneTileIOdef oneTileIO (output_ptr input_ptr : RegionName) (N TILE_N : Nat) :
MaskedTileKernelIO₁ where
kernel := softmax_kernel_online_v2_one_tile output_ptr input_ptr N TILE_N
inp := input_ptr
out := output_ptr
shape := [TILE_N]
read := fun pid idx => pid * N + idx.1.val
write := fun pid idx => pid * N + idx.1.val
mask := fun _pid idx => idx.1.val < N
Referenced model/specification definitions: Python source |
|---|
softmax_reducevsoftmax_reducev_genuine_output_compute_correct_generalMathematical execution | Configured model / stage Modeled softmax-weighted V reduction with supplied token/index metadata and positive block/stream premises. The host metadata pipeline and concrete floating behavior remain external. | Statement and preconditionsspecification softmax_reducev_genuine_output_compute_correct_general
(Logics V Out : RegionName) (BLoc : Region .int) (BStartLoc BSeqLen : Region .nat)
(mil slh slb svbs svh svd sob soh sod sb ss BLOCK_DMODEL BLOCK_N : Nat)
(other_kv_index : Int)
(hD : 0 < BLOCK_DMODEL) (hN : 0 < BLOCK_N)
(s : BlockState) (hundef : ∀ rg o, s.undef rg o = 0)
(hseqmod : srSeqLen s BSeqLen.cast % BLOCK_N = 0) (hseqpos : 0 < srSeqLen s BSeqLen.cast)
(hOutInj : ∀ s0 : BlockState, Function.Injective (fun i : Fin BLOCK_DMODEL => outOffsetG s0 sob soh sod i))
(mr : ℝ)
(hM : srRunningMax (srQkFG s Logics BStartLoc.cast BSeqLen.cast slh slb)
(srVFG s V BLoc BSeqLen.cast mil sb ss svbs svh svd BLOCK_DMODEL)
(srSeqLen s BSeqLen.cast) (⟨0, hD⟩ : Fin BLOCK_DMODEL) = (mr : WithBot ℝ)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := softmax_reducev_surface Logics V Out BLoc BStartLoc BSeqLen
mil slh slb svbs svh svd sob soh sod sb ss BLOCK_DMODEL BLOCK_N other_kv_index)
(initialState := s)
(write := fun d : Fin BLOCK_DMODEL => some (Out, outOffsetG s sob soh sod d))
(expected := fun d : Fin BLOCK_DMODEL =>
softmaxReducevWeightedSum (srQkFG s Logics BStartLoc.cast BSeqLen.cast slh slb) mr
(srVFG s V BLoc BSeqLen.cast mil sb ss svbs svh svd BLOCK_DMODEL) d)
Referenced model/specification definitions: Python source |
|---|
softmax_reducevsoftmax_reducev_io_correctnessAbstract cast/store rounding | Configured model / stage Modeled softmax-weighted V reduction with supplied token/index metadata and positive block/stream premises. The host metadata pipeline and concrete floating behavior remain external. | Statement and preconditionsspecification softmax_reducev_io_correctness (R : RoundingModel)
(Logics V Out : RegionName) (BLoc : Region .int)
(BStartLoc BSeqLen : Region .nat)
(mil slh slb svbs svh svd sob soh sod sb ss BLOCK_DMODEL BLOCK_N : Nat)
(other_kv_index : Int) (hD : 0 < BLOCK_DMODEL) (hN : 0 < BLOCK_N)
(hOInj : ∀ pid₀ pid₁ : Nat, Function.Injective
(fun i : Fin BLOCK_DMODEL => pid₀ * sob + pid₁ * soh + i.val * sod)) :
softmaxReducevIO Logics V Out BLoc BStartLoc BSeqLen mil slh slb svbs svh svd
sob soh sod sb ss BLOCK_DMODEL BLOCK_N other_kv_index ⊨[R]
fun _ _ _ m xs ys j =>
softmaxReducevIOSpec BLOCK_N BLOCK_DMODEL (srIOT mil BLOCK_N)
(m (⟨0, by omega⟩ : Fin 2)) hN xs ys j
IO definition: softmaxReducevIOdef softmaxReducevIO (Logics V Out : RegionName) (BLoc : Region .int)
(BStartLoc BSeqLen : Region .nat)
(mil slh slb svbs svh svd sob soh sod sb ss BLOCK_DMODEL BLOCK_N : Nat)
(other_kv_index : Int) : StreamMetaGatherMasked3DKernelIO₂ where
kernel := softmax_reducev_surface Logics V Out BLoc BStartLoc BSeqLen
mil slh slb svbs svh svd sob soh sod sb ss BLOCK_DMODEL BLOCK_N other_kv_index
inp1 := Logics
inp2 := V
out := Out
nMeta := 2
sty := fun _ => ChanTy.nat
mbuf := srIOMetaBuf BSeqLen BStartLoc
mwin := fun _ pid₀ _ _ => pid₀
gbuf := BLoc.cast
gty := ChanTy.int
Bg := BLOCK_N
gother := other_kv_index
T := srIOT mil BLOCK_N
B1 := BLOCK_N
B2 := BLOCK_N * BLOCK_DMODEL
C := BLOCK_DMODEL
outDType := .real
pre := fun _ _ _ m =>
m (⟨0, by omega⟩ : Fin 2) ≤ mil ∧ 0 < m (⟨0, by omega⟩ : Fin 2)
∧ m (⟨0, by omega⟩ : Fin 2) % BLOCK_N = 0
gread := fun pid₀ _ _ m t jL =>
pid₀ * sb + (mil - m (⟨0, by omega⟩ : Fin 2)) * ss
+ (t.val * BLOCK_N + jL.val) * ss
gmask := fun _ _ _ m t jL =>
t.val * BLOCK_N + jL.val < m (⟨0, by omega⟩ : Fin 2)
read1 := fun _ pid₁ _ m t jL =>
pid₁ * slh + (m (⟨1, by omega⟩ : Fin 2) + (t.val * BLOCK_N + jL.val)) * slb
mask1 := fun _ _ _ m t jL =>
t.val * BLOCK_N + jL.val < m (⟨0, by omega⟩ : Fin 2)
read2 := fun _ pid₁ _ _ G t j =>
(G t (Lane2D.decode j).1 * svbs).toNat + pid₁ * svh
+ (Lane2D.decode j).2.1.val * svd
mask2 := fun _ _ _ _ _ _ => True
write := fun pid₀ pid₁ _ _ i => pid₀ * sob + pid₁ * soh + i.val * sod
writeMask := fun _ _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
softmax_triton1softmax_kernel_correctnessMathematical execution | Configured model / stage One masked row softmax with positive BLOCK_SIZE and declared input/output strides; no host multi-kernel composition or concrete exp/rounding guarantee. | Statement and preconditionsspecification softmax_kernel_correctness
(output_ptr input_ptr : RegionName)
(input_row_stride output_row_stride n_cols BLOCK_SIZE : Nat)
(hB : 0 < BLOCK_SIZE) :
softmaxIO output_ptr input_ptr input_row_stride output_row_stride
n_cols BLOCK_SIZE ⊨
fun xs i => softmaxSpec n_cols BLOCK_SIZE xs i
IO definition: softmaxIOdef softmaxIO (output_ptr input_ptr : RegionName)
(input_row_stride output_row_stride n_cols BLOCK_SIZE : Nat) :
MaskedKernelIO₁ where
kernel := softmax_kernel output_ptr input_ptr input_row_stride
output_row_stride n_cols BLOCK_SIZE
inp := input_ptr
out := output_ptr
B := BLOCK_SIZE
read := fun pid => pid * input_row_stride
write := fun pid => pid * output_row_stride
mask := fun _ j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
softmax_triton2softmax_kernel_correctnessMathematical execution | Configured model / stage One masked row softmax with positive BLOCK_SIZE and explicit row strides, under mathematical exp/division semantics. | Statement and preconditionsspecification softmax_kernel_correctness
(output_ptr input_ptr : RegionName)
(input_row_stride output_row_stride n_cols BLOCK_SIZE : Nat)
(hB : 0 < BLOCK_SIZE) :
softmaxIO output_ptr input_ptr input_row_stride output_row_stride
n_cols BLOCK_SIZE ⊨
fun xs i => softmaxSpec n_cols BLOCK_SIZE xs i
IO definition: softmaxIOdef softmaxIO (output_ptr input_ptr : RegionName)
(input_row_stride output_row_stride n_cols BLOCK_SIZE : Nat) :
MaskedKernelIO₁ where
kernel := softmax_kernel output_ptr input_ptr input_row_stride
output_row_stride n_cols BLOCK_SIZE
inp := input_ptr
out := output_ptr
B := BLOCK_SIZE
read := fun pid => pid * input_row_stride
write := fun pid => pid * output_row_stride
mask := fun _ j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
softmax_triton3softmax_kernel_correctnessMathematical execution | Configured model / stage Separate USE_MASK=false and USE_MASK=true row contracts with positive block size. Only the corresponding input/mask windows and modeled arithmetic are claimed. | Statement and preconditionsspecification softmax_kernel_correctness
(output_ptr input_ptr mask_ptr : RegionName)
(row_stride n_cols BLOCK_SIZE : Nat)
(hB : 0 < BLOCK_SIZE) :
softmaxIO output_ptr input_ptr mask_ptr row_stride n_cols BLOCK_SIZE ⊨
fun xs i => softmaxSpec n_cols BLOCK_SIZE none xs i
IO definition: softmaxIOdef softmaxIO (output_ptr input_ptr mask_ptr : RegionName)
(row_stride n_cols BLOCK_SIZE : Nat) : MaskedKernelIO₁ where
kernel := softmax_kernel output_ptr input_ptr mask_ptr row_stride n_cols
BLOCK_SIZE Bool.false
inp := input_ptr
out := output_ptr
B := BLOCK_SIZE
read := fun pid => pid * row_stride
write := fun pid => pid * row_stride
mask := fun _ j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
softmax_triton3softmax_kernel_masked_correctnessMathematical execution | Configured model / stage Separate USE_MASK=false and USE_MASK=true row contracts with positive block size. Only the corresponding input/mask windows and modeled arithmetic are claimed. | Statement and preconditionsspecification softmax_kernel_masked_correctness
(output_ptr input_ptr mask_ptr : RegionName)
(row_stride n_cols BLOCK_SIZE : Nat)
(hB : 0 < BLOCK_SIZE) :
softmaxMaskedIO output_ptr input_ptr mask_ptr row_stride n_cols BLOCK_SIZE ⊨
fun xs ms i => softmaxSpec n_cols BLOCK_SIZE (some ms) xs i
IO definition: softmaxMaskedIOdef softmaxMaskedIO (output_ptr input_ptr mask_ptr : RegionName)
(row_stride n_cols BLOCK_SIZE : Nat) : MaskedKernelIO₂ where
kernel := softmax_kernel output_ptr input_ptr mask_ptr row_stride n_cols
BLOCK_SIZE Bool.true
in1 := input_ptr
in2 := mask_ptr
out := output_ptr
B := BLOCK_SIZE
read1 := fun pid => pid * row_stride
read2 := fun pid => pid * row_stride
write := fun pid => pid * row_stride
mask := fun _ j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
square_matrixsquare_kernel_correctnessMathematical execution | Configured model / stage Per-program masked elementwise square with explicit row strides; output window bounds and separation come from the IO contract. | Statement and preconditionsspecification square_kernel_correctness
(output_ptr input_ptr : RegionName)
(input_row_stride output_row_stride n_cols BLOCK_SIZE : Nat) :
squareIO output_ptr input_ptr input_row_stride output_row_stride
n_cols BLOCK_SIZE ⊨
fun xs i => xs i * xs i
IO definition: squareIOdef squareIO (output_ptr input_ptr : RegionName)
(input_row_stride output_row_stride n_cols BLOCK_SIZE : Nat) :
MaskedKernelIO₁ where
kernel := square_kernel output_ptr input_ptr input_row_stride
output_row_stride n_cols BLOCK_SIZE
inp := input_ptr
out := output_ptr
B := BLOCK_SIZE
read := fun pid => pid * input_row_stride
write := fun pid => pid * output_row_stride
mask := fun _ j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
swiglu_backwardswiglu_bwd_kernel_correctnessMathematical execution | Configured model / stage Backward DX/DY and optionally recomputed OUT under the declared RECOMPUTE_OUTPUT flag and region/window premises. Mathematical sigmoid is not a hardware approximation certificate. | Statement and preconditionsspecification swiglu_bwd_kernel_correctness
(X Y DOUT OUT DX DY : RegionName)
(stride_x_row stride_y_row stride_dout_row stride_out_row
stride_dx_row stride_dy_row ncols BLOCK_N : Nat)
(RECOMPUTE_OUTPUT : Bool)
(hDXDY : DX ≠ DY) (hOUTDX : OUT ≠ DX) (hOUTDY : OUT ≠ DY) :
swigluBackwardIO X Y DOUT OUT DX DY stride_x_row stride_y_row
stride_dout_row stride_out_row stride_dx_row stride_dy_row ncols
BLOCK_N RECOMPUTE_OUTPUT ⊨
fun _ _ xs ys douts =>
(fun j => TiledActivation.swigluBwdA (douts j) (xs j) (ys j),
fun j => TiledActivation.swigluBwdB (douts j) (xs j),
fun j => TiledActivation.swiglu (xs j) (ys j))
IO definition: swigluBackwardIOdef swigluBackwardIO (X Y DOUT OUT DX DY : RegionName)
(stride_x_row stride_y_row stride_dout_row stride_out_row
stride_dx_row stride_dy_row ncols BLOCK_N : Nat)
(RECOMPUTE_OUTPUT : Bool) : Masked2DKernelIO₃ₓ₃ where
kernel := swiglu_bwd_kernel X Y DOUT OUT DX DY
stride_x_row stride_y_row stride_dout_row stride_out_row
stride_dx_row stride_dy_row ncols BLOCK_N RECOMPUTE_OUTPUT
in1 := X
in2 := Y
in3 := DOUT
out1 := DX
out2 := DY
out3 := OUT
B := BLOCK_N
read1 := fun pid₀ pid₁ j => pid₀ * stride_x_row + (pid₁ * BLOCK_N + j.val)
read2 := fun pid₀ pid₁ j => pid₀ * stride_y_row + (pid₁ * BLOCK_N + j.val)
read3 := fun pid₀ pid₁ j => pid₀ * stride_dout_row + (pid₁ * BLOCK_N + j.val)
write1 := fun pid₀ pid₁ j => pid₀ * stride_dx_row + (pid₁ * BLOCK_N + j.val)
write2 := fun pid₀ pid₁ j => pid₀ * stride_dy_row + (pid₁ * BLOCK_N + j.val)
write3 := fun pid₀ pid₁ j => pid₀ * stride_out_row + (pid₁ * BLOCK_N + j.val)
mask := fun _ pid₁ j => pid₁ * BLOCK_N + j.val < ncols
writeMask3 := fun _ pid₁ j =>
RECOMPUTE_OUTPUT = Bool.true ∧ pid₁ * BLOCK_N + j.val < ncols
Referenced model/specification definitions: Python source |
|---|
swiglu_fwdswiglu_fwd_kernel_correctnessMathematical execution | Configured model / stage Per-row SwiGLU forward on the declared input/output strides and masks, using mathematical sigmoid and multiplication. | Statement and preconditionsspecification swiglu_fwd_kernel_correctness
(X Y OUT : RegionName)
(stride_x_row stride_y_row stride_out_row ncols BLOCK_N : Nat) :
swigluIO X Y OUT stride_x_row stride_y_row stride_out_row ncols BLOCK_N
⊨ fun _ _ xs ys i => TiledActivation.swiglu (xs i) (ys i)
IO definition: swigluIOdef swigluIO (X Y OUT : RegionName)
(stride_x_row stride_y_row stride_out_row ncols BLOCK_N : Nat) :
Masked2DKernelIO₂ where
kernel := swiglu_fwd_kernel X Y OUT stride_x_row stride_y_row stride_out_row
ncols BLOCK_N
in1 := X
in2 := Y
out := OUT
B := BLOCK_N
read1 := fun pid₀ pid₁ j => pid₀ * stride_x_row + pid₁ * BLOCK_N + j.val
read2 := fun pid₀ pid₁ j => pid₀ * stride_y_row + pid₁ * BLOCK_N + j.val
write := fun pid₀ pid₁ j => pid₀ * stride_out_row + pid₁ * BLOCK_N + j.val
mask := fun _ pid₁ j => pid₁ * BLOCK_N + j.val < ncols
Referenced model/specification definitions: Python source |
|---|
swiglu_tritonswiglu_forward_kernel_correctnessMathematical execution | Configured model / stage Separate forward and in-place backward SwiGLU contracts. The backward theorem requires distinct A/B buffers; no host autograd or hardware transcendental guarantee. | Statement and preconditionsspecification swiglu_forward_kernel_correctness
(A B C : RegionName)
(stride n_cols BLOCK_SIZE : Nat) :
swigluFwdIO A B C stride n_cols BLOCK_SIZE
⊨ fun as bs i => TiledActivation.swiglu (as i) (bs i)
IO definition: swigluFwdIOdef swigluFwdIO (A B C : RegionName)
(stride n_cols BLOCK_SIZE : Nat) : MaskedKernelIO₂ where
kernel := swiglu_forward_kernel A B C stride n_cols BLOCK_SIZE
in1 := A
in2 := B
out := C
B := BLOCK_SIZE
read1 := fun pid => pid * stride
read2 := fun pid => pid * stride
write := fun pid => pid * stride
mask := fun _ j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
swiglu_tritonswiglu_backward_kernel_correctnessMathematical execution | Configured model / stage Separate forward and in-place backward SwiGLU contracts. The backward theorem requires distinct A/B buffers; no host autograd or hardware transcendental guarantee. | Statement and preconditionsspecification swiglu_backward_kernel_correctness
(DC A B : RegionName)
(stride n_cols BLOCK_SIZE : Nat)
(hAB : A ≠ B) :
swigluBwdIO DC A B stride n_cols BLOCK_SIZE
⊨ fun dcs as bs =>
(fun i => TiledActivation.swigluBwdA (dcs i) (as i) (bs i),
fun i => TiledActivation.swigluBwdB (dcs i) (as i))
IO definition: swigluBwdIOdef swigluBwdIO (DC A B : RegionName)
(stride n_cols BLOCK_SIZE : Nat) : MaskedKernelIO₃ₓ₂ where
kernel := swiglu_backward_kernel DC A B stride n_cols BLOCK_SIZE
bufs := [DC, A, B] -- A and B are updated in place
in1 := DC
in2 := A
in3 := B
out1 := A -- = in2: `da` overwrites the gate input in place
out2 := B -- = in3: `db` overwrites the value input in place
B := BLOCK_SIZE
read1 := fun pid => pid * stride
read2 := fun pid => pid * stride
read3 := fun pid => pid * stride
write1 := fun pid => pid * stride
write2 := fun pid => pid * stride
mask := fun _ j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
token_attn_llama2token_attn_llama2_output_summary_generalMathematical execution | Configured model / stage Token Q/K score stage with supplied token-location and sequence metadata. It does not compute the later softmax or V-reduction stages. | Statement and preconditionsspecification token_attn_llama2_output_summary_general
(Q K : RegionName) (sm_scale : ℝ)
(B_Loc B_Start_Loc B_Seqlen : Region .nat) (Att_Out : RegionName)
(max_input_len stride_b_loc_b stride_b_loc_s stride_qbs stride_qh stride_qd
stride_kbs stride_kh stride_kd att_stride_h att_stride_bs kv_group_num
BLOCK_DMODEL BLOCK_N : Nat)
(s : BlockState) (hundef : ∀ rg o, s.undef rg o = 0)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_N => outOffset s B_Start_Loc att_stride_h att_stride_bs BLOCK_N i)) :
(∃ alg, (token_attn_llama2_surface Q K sm_scale B_Loc B_Start_Loc B_Seqlen
Att_Out max_input_len stride_b_loc_b stride_b_loc_s stride_qbs stride_qh
stride_qd stride_kbs stride_kh stride_kd att_stride_h att_stride_bs
kv_group_num BLOCK_DMODEL BLOCK_N).toAlgorithm? = Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := token_attn_llama2_surface Q K sm_scale B_Loc B_Start_Loc
B_Seqlen Att_Out max_input_len stride_b_loc_b stride_b_loc_s stride_qbs
stride_qh stride_qd stride_kbs stride_kh stride_kd att_stride_h
att_stride_bs kv_group_num BLOCK_DMODEL BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_N => active s B_Seqlen max_input_len BLOCK_N i)
(fun i => (Att_Out, outOffset s B_Start_Loc att_stride_h att_stride_bs BLOCK_N i)))
(expected := fun i : Fin BLOCK_N =>
tokenAttnLlama2ClosedForm s Q K sm_scale B_Loc B_Start_Loc B_Seqlen Att_Out
max_input_len stride_b_loc_b stride_b_loc_s stride_qbs stride_qh stride_qd
stride_kbs stride_kh stride_kd att_stride_h att_stride_bs kv_group_num
BLOCK_DMODEL BLOCK_N i))
Referenced model/specification definitions: Python source |
|---|
token_attn_llama2token_attn_llama2_io_correctnessAbstract cast/store rounding | Configured model / stage Token Q/K score stage with supplied token-location and sequence metadata. It does not compute the later softmax or V-reduction stages. | Statement and preconditionsspecification token_attn_llama2_io_correctness (R : RoundingModel)
(Q K : RegionName) (sm_scale : ℝ) (B_Loc B_Start_Loc B_Seqlen : Region .nat)
(Att_Out : RegionName)
(mil sblb sbls sqbs sqh sqd skbs skh skd ash asbs kvg BD BN : Nat)
(hOutInj : ∀ (pid₁ pid₂ base : Nat), Function.Injective
(fun i : Fin BN => pid₁ * ash + (base + (pid₂ * BN + i.val)) * asbs)) :
llama2IO Q K sm_scale B_Loc B_Start_Loc B_Seqlen Att_Out mil sblb sbls sqbs sqh sqd
skbs skh skd ash asbs kvg BD BN ⊨[R]
fun _ _ _ _ xs ys i => llama2IOSpec BD BN sm_scale xs ys i
IO definition: llama2IOdef llama2IO (Q K : RegionName) (sm_scale : ℝ)
(B_Loc B_Start_Loc B_Seqlen : Region .nat) (Att_Out : RegionName)
(mil sblb sbls sqbs sqh sqd skbs skh skd ash asbs kvg BD BN : Nat) :
StreamMetaGatherMasked3DKernelIO₂ where
kernel := token_attn_llama2_surface Q K sm_scale B_Loc B_Start_Loc B_Seqlen Att_Out
mil sblb sbls sqbs sqh sqd skbs skh skd ash asbs kvg BD BN
inp1 := Q
inp2 := K
out := Att_Out
nMeta := 2
sty := fun _ => ChanTy.nat
mbuf := llama2IOMetaBuf B_Start_Loc B_Seqlen
mwin := fun _ pid₀ _ _ => pid₀
gbuf := B_Loc.cast
gty := ChanTy.nat
Bg := BN
gother := 0
T := 1
B1 := BD
B2 := BN * BD
C := BN
outDType := .real
pre := fun _ _ _ _ => True
gread := fun pid₀ _ pid₂ m _ j =>
sblb * pid₀ + sbls * ((mil - m (⟨0, by omega⟩ : Fin 2)) + (pid₂ * BN + j.val))
gmask := fun _ _ pid₂ m _ j =>
(mil - m (⟨0, by omega⟩ : Fin 2)) + (pid₂ * BN + j.val) < mil
read1 := fun pid₀ pid₁ _ _ t d => pid₀ * sqbs + pid₁ * sqh + d.val * sqd + t.val
mask1 := fun _ _ _ _ _ _ => True
read2 := fun _ pid₁ _ _ G t j =>
G t (Lane2D.decode j).1 * skbs + (pid₁ / kvg) * skh + (Lane2D.decode j).2.1.val * skd
mask2 := fun _ _ pid₂ m _ j =>
(mil - m (⟨0, by omega⟩ : Fin 2)) + (pid₂ * BN + (Lane2D.decode j).1.val) < mil
write := fun _ pid₁ pid₂ m i =>
pid₁ * ash + (m (⟨1, by omega⟩ : Fin 2) + (pid₂ * BN + i.val)) * asbs
writeMask := fun _ _ pid₂ m i =>
pid₂ * BN < m (⟨0, by omega⟩ : Fin 2) ∧
(mil - m (⟨0, by omega⟩ : Fin 2)) + (pid₂ * BN + i.val) < mil
Referenced model/specification definitions: Python source |
|---|
token_attn_mistraltoken_attn_mistral_output_summary_generalMathematical execution | Configured model / stage Sliding-window weighted V reduction consumes supplied probabilities and token metadata. The score/softmax producers are outside this original stage contract. | Statement and preconditionsspecification token_attn_mistral_output_summary_general
(Prob V Out : RegionName)
(Req_to_tokens B_req_idx : Region .nat) (B_Start_Loc : RegionName)
(B_Seqlen B_Att_Start_Loc B_Att_Seqlen : Region .nat)
(stride_req_to_tokens_b stride_req_to_tokens_s stride_ph stride_pbs
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
kv_group_num sliding_window BLOCK_DMODEL BLOCK_N : Nat)
(hpbs : stride_pbs = 1) (hrts : stride_req_to_tokens_s = 1) (hBN : 0 < BLOCK_N)
(s : BlockState) (hundef : ∀ rg o, s.undef rg o = 0)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_DMODEL => outOffset s stride_obs stride_oh stride_od i)) :
(∃ alg, (token_attn_mistral_surface Prob V Out Req_to_tokens B_req_idx
B_Start_Loc B_Seqlen B_Att_Start_Loc B_Att_Seqlen
stride_req_to_tokens_b stride_req_to_tokens_s stride_ph stride_pbs
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
kv_group_num sliding_window BLOCK_DMODEL BLOCK_N).toAlgorithm? =
Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := token_attn_mistral_surface Prob V Out Req_to_tokens B_req_idx
B_Start_Loc B_Seqlen B_Att_Start_Loc B_Att_Seqlen
stride_req_to_tokens_b stride_req_to_tokens_s stride_ph stride_pbs
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
kv_group_num sliding_window BLOCK_DMODEL BLOCK_N)
(initialState := s)
(write := fun i : Fin BLOCK_DMODEL =>
some (Out, outOffset s stride_obs stride_oh stride_od i))
(expected := fun i : Fin BLOCK_DMODEL =>
tokenAttnMistralClosedForm s Prob V Req_to_tokens B_req_idx B_Att_Start_Loc
B_Seqlen B_Att_Seqlen stride_req_to_tokens_b stride_req_to_tokens_s
stride_ph stride_pbs stride_vbs stride_vh stride_vd kv_group_num
sliding_window BLOCK_DMODEL i))
Referenced model/specification definitions: Python source |
|---|
token_attn_mistraltoken_attn_mistral_io_correctnessAbstract cast/store rounding | Configured model / stage Sliding-window weighted V reduction consumes supplied probabilities and token metadata. The score/softmax producers are outside this original stage contract. | Statement and preconditionsspecification token_attn_mistral_io_correctness (R : RoundingModel)
(Prob V Out : RegionName) (Req_to_tokens B_req_idx : Region .nat)
(B_Start_Loc : RegionName) (B_Seqlen B_Att_Start_Loc B_Att_Seqlen : Region .nat)
(srtb srts sph spbs svbs svh svd sobs soh sod kvg sw BLOCK_DMODEL BLOCK_N T : Nat)
(hpbs : spbs = 1) (hrts : srts = 1) (hBN : 0 < BLOCK_N)
(hOutInj : ∀ pid₀ pid₁ : Nat, Function.Injective
(fun i : Fin BLOCK_DMODEL => pid₀ * sobs + pid₁ * soh + i.val * sod)) :
mistralIO Prob V Out Req_to_tokens B_req_idx B_Start_Loc B_Seqlen B_Att_Start_Loc
B_Att_Seqlen srtb srts sph spbs svbs svh svd sobs soh sod kvg sw BLOCK_DMODEL
BLOCK_N T ⊨[R]
fun _ _ _ m xs ys j =>
mistralIOSpec BLOCK_N BLOCK_DMODEL T (m (⟨3, by omega⟩ : Fin 4))
(m (⟨0, by omega⟩ : Fin 4)) sw hBN xs ys j
IO definition: mistralIOdef mistralIO (Prob V Out : RegionName)
(Req_to_tokens B_req_idx : Region .nat) (B_Start_Loc : RegionName)
(B_Seqlen B_Att_Start_Loc B_Att_Seqlen : Region .nat)
(srtb srts sph spbs svbs svh svd sobs soh sod kvg sw BLOCK_DMODEL BLOCK_N T : Nat) :
StreamMetaGatherMasked3DKernelIO₂ where
kernel := token_attn_mistral_surface Prob V Out Req_to_tokens B_req_idx B_Start_Loc
B_Seqlen B_Att_Start_Loc B_Att_Seqlen srtb srts sph spbs svbs svh svd sobs soh sod
kvg sw BLOCK_DMODEL BLOCK_N
inp1 := Prob
inp2 := V
out := Out
nMeta := 4
sty := fun _ => ChanTy.nat
mbuf := mistralIOMetaBuf B_Seqlen B_Att_Start_Loc B_req_idx B_Att_Seqlen
mwin := fun _ pid₀ _ _ => pid₀
gbuf := Req_to_tokens.cast
gty := ChanTy.nat
Bg := BLOCK_N
gother := 0
T := T
B1 := BLOCK_N
B2 := BLOCK_N * BLOCK_DMODEL
C := BLOCK_DMODEL
outDType := .real
pre := fun _ _ _ m => m (⟨3, by omega⟩ : Fin 4) ≤ T * BLOCK_N
gread := fun _ _ _ m t jL =>
m (⟨2, by omega⟩ : Fin 4) * srtb
+ ((m (⟨0, by omega⟩ : Fin 4) - sw) + (t.val * BLOCK_N + jL.val)) * srts
gmask := fun _ _ _ m t jL =>
(m (⟨0, by omega⟩ : Fin 4) - sw) + (t.val * BLOCK_N + jL.val)
< m (⟨0, by omega⟩ : Fin 4)
read1 := fun _ pid₁ _ m t jL =>
pid₁ * sph + (m (⟨1, by omega⟩ : Fin 4) + (t.val * BLOCK_N + jL.val)) * spbs
mask1 := fun _ _ _ m t jL => t.val * BLOCK_N + jL.val < m (⟨3, by omega⟩ : Fin 4)
read2 := fun _ pid₁ _ _ G t j =>
G t (Lane2D.decode j).1 * svbs + (pid₁ / kvg) * svh
+ (Lane2D.decode j).2.1.val * svd
mask2 := fun _ _ _ m t j =>
(m (⟨0, by omega⟩ : Fin 4) - sw) + (t.val * BLOCK_N + (Lane2D.decode j).1.val)
< m (⟨0, by omega⟩ : Fin 4)
write := fun pid₀ pid₁ _ _ i => pid₀ * sobs + pid₁ * soh + i.val * sod
writeMask := fun _ _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
token_attn_reduceVtoken_attn_reducev_output_summary_generalMathematical execution | Configured model / stage Weighted V reduction consumes supplied Prob and index metadata. The separate score and softmax stages and host composition are not included. | Statement and preconditionsspecification token_attn_reducev_output_summary_general
(Prob V Out : RegionName)
(Req_to_tokens B_req_idx B_Start_Loc B_Seqlen : Region .nat)
(stride_req_to_tokens_b stride_req_to_tokens_s stride_ph stride_pbs
stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
kv_group_num BLOCK_DMODEL BLOCK_N : Nat)
(hBD : 0 < BLOCK_DMODEL) (hBN : 0 < BLOCK_N)
(hpbs : stride_pbs = 1) (hrts : stride_req_to_tokens_s = 1)
(s : BlockState) (hundef : ∀ rg o, s.undef rg o = 0)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_DMODEL => outOffset s stride_obs stride_oh stride_od i)) :
(∃ alg, (token_attn_reducev_surface Prob V Out Req_to_tokens B_req_idx
B_Start_Loc B_Seqlen stride_req_to_tokens_b stride_req_to_tokens_s stride_ph
stride_pbs stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
kv_group_num BLOCK_DMODEL BLOCK_N).toAlgorithm? = Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := token_attn_reducev_surface Prob V Out Req_to_tokens B_req_idx
B_Start_Loc B_Seqlen stride_req_to_tokens_b stride_req_to_tokens_s stride_ph
stride_pbs stride_vbs stride_vh stride_vd stride_obs stride_oh stride_od
kv_group_num BLOCK_DMODEL BLOCK_N)
(initialState := s)
(write := fun i : Fin BLOCK_DMODEL =>
some (Out, outOffset s stride_obs stride_oh stride_od i))
(expected := fun i : Fin BLOCK_DMODEL =>
tokenAttnReduceVClosedForm s Prob V Req_to_tokens B_req_idx
B_Start_Loc B_Seqlen stride_req_to_tokens_b stride_req_to_tokens_s stride_ph
stride_pbs stride_vbs stride_vh stride_vd kv_group_num BLOCK_DMODEL i))
Referenced model/specification definitions: Python source |
|---|
token_attn_reduceVtoken_attn_reducev_io_correctnessAbstract cast/store rounding | Configured model / stage Weighted V reduction consumes supplied Prob and index metadata. The separate score and softmax stages and host composition are not included. | Statement and preconditionsspecification token_attn_reducev_io_correctness (R : RoundingModel)
(Prob V Out : RegionName)
(Req_to_tokens B_req_idx B_Start_Loc B_Seqlen : Region .nat)
(srtb srts sph spbs svbs svh svd sobs soh sod kvg BLOCK_DMODEL BLOCK_N T : Nat)
(hpbs : spbs = 1) (hrts : srts = 1) (hBN : 0 < BLOCK_N)
(hOutInj : ∀ pid₀ pid₁ : Nat, Function.Injective
(fun i : Fin BLOCK_DMODEL => pid₀ * sobs + pid₁ * soh + i.val * sod)) :
tokenAttnReduceVIO Prob V Out Req_to_tokens B_req_idx B_Start_Loc B_Seqlen srtb srts
sph spbs svbs svh svd sobs soh sod kvg BLOCK_DMODEL BLOCK_N T ⊨[R]
fun _ _ _ m xs ys j =>
tokenAttnReduceVIOSpec BLOCK_N BLOCK_DMODEL T (m (⟨0, by omega⟩ : Fin 3))
hBN xs ys j
IO definition: tokenAttnReduceVIOdef tokenAttnReduceVIO (Prob V Out : RegionName)
(Req_to_tokens B_req_idx B_Start_Loc B_Seqlen : Region .nat)
(srtb srts sph spbs svbs svh svd sobs soh sod kvg BLOCK_DMODEL BLOCK_N T : Nat) :
StreamMetaGatherMasked3DKernelIO₂ where
kernel := token_attn_reducev_surface Prob V Out Req_to_tokens B_req_idx B_Start_Loc
B_Seqlen srtb srts sph spbs svbs svh svd sobs soh sod kvg BLOCK_DMODEL BLOCK_N
inp1 := Prob
inp2 := V
out := Out
nMeta := 3
sty := fun _ => ChanTy.nat
mbuf := rvIOMetaBuf B_Seqlen B_Start_Loc B_req_idx
mwin := fun _ pid₀ _ _ => pid₀
gbuf := Req_to_tokens.cast
gty := ChanTy.nat
Bg := BLOCK_N
gother := 0
T := T
B1 := BLOCK_N
B2 := BLOCK_N * BLOCK_DMODEL
C := BLOCK_DMODEL
outDType := .real
pre := fun _ _ _ m => m (⟨0, by omega⟩ : Fin 3) ≤ T * BLOCK_N
gread := fun _ _ _ m t jL =>
m (⟨2, by omega⟩ : Fin 3) * srtb + (t.val * BLOCK_N + jL.val) * srts
gmask := fun _ _ _ m t jL => t.val * BLOCK_N + jL.val < m (⟨0, by omega⟩ : Fin 3)
read1 := fun _ pid₁ _ m t jL =>
pid₁ * sph + (m (⟨1, by omega⟩ : Fin 3) + (t.val * BLOCK_N + jL.val)) * spbs
mask1 := fun _ _ _ m t jL => t.val * BLOCK_N + jL.val < m (⟨0, by omega⟩ : Fin 3)
read2 := fun _ pid₁ _ _ G t j =>
G t (Lane2D.decode j).1 * svbs + (pid₁ / kvg) * svh
+ (Lane2D.decode j).2.1.val * svd
mask2 := fun _ _ _ m t j =>
t.val * BLOCK_N + (Lane2D.decode j).1.val < m (⟨0, by omega⟩ : Fin 3)
write := fun pid₀ pid₁ _ _ i => pid₀ * sobs + pid₁ * soh + i.val * sod
writeMask := fun _ _ _ _ _ => True
Referenced model/specification definitions: Python source |
|---|
token_softmax_bloomtoken_softmax_bloom_correctnessMathematical execution | Configured model / stage Per-sequence token softmax consumes supplied logits and start/length metadata. IO bounds and injective probability addresses are prerequisites. | Statement and preconditionsspecification token_softmax_bloom_correctness
(Logics B_Start_Loc B_Seqlen Prob_Out : RegionName)
(stride_logic_h stride_logic_bs stride_prob_h stride_prob_bs BLOCK_SIZE : Nat)
(hB : 0 < BLOCK_SIZE)
(hOutInj : ∀ pid₁ m₁, Function.Injective
(fun j : Fin BLOCK_SIZE =>
pid₁ * stride_prob_h + (m₁ + j.val) * stride_prob_bs)) :
tokenSoftmaxIO Logics B_Start_Loc B_Seqlen Prob_Out
stride_logic_h stride_logic_bs stride_prob_h stride_prob_bs BLOCK_SIZE
⊨ fun _ _ _ m₂ xs j => tokenSoftmaxSpecPure m₂ xs j
IO definition: tokenSoftmaxIOdef tokenSoftmaxIO (Logics B_Start_Loc B_Seqlen Prob_Out : RegionName)
(stride_logic_h stride_logic_bs stride_prob_h stride_prob_bs
BLOCK_SIZE : Nat) : MetaMasked2DKernelIO₁ where
kernel := token_softmax_surface Logics B_Start_Loc B_Seqlen Prob_Out
stride_logic_h stride_logic_bs stride_prob_h stride_prob_bs BLOCK_SIZE
mbuf1 := B_Start_Loc
mbuf2 := B_Seqlen
inp := Logics
out := Prob_Out
B := BLOCK_SIZE
mwin1 := fun pid₀ _ => pid₀
mwin2 := fun pid₀ _ => pid₀
read := fun _ pid₁ m₁ _ j =>
pid₁ * stride_logic_h + (m₁ + j.val) * stride_logic_bs
write := fun _ pid₁ m₁ _ j =>
pid₁ * stride_prob_h + (m₁ + j.val) * stride_prob_bs
mask := fun _ _ _ m₂ j => j.val < m₂
Referenced model/specification definitions: Python source |
|---|
token_softmax_llamatoken_softmax_llama_correctnessMathematical execution | Configured model / stage Per-sequence token softmax with supplied logits/start/length metadata and explicit output-address injectivity; surrounding attention stages are external. | Statement and preconditionsspecification token_softmax_llama_correctness
(Logics B_Start_Loc B_Seqlen Prob_Out : RegionName)
(stride_logic_h stride_logic_bs stride_prob_h stride_prob_bs
BLOCK_SIZE : Nat)
(hB : 0 < BLOCK_SIZE)
(hOutInj : ∀ pid₁ m₁ : Nat, Function.Injective
(fun i : Fin BLOCK_SIZE =>
pid₁ * stride_prob_h + (m₁ + i.val) * stride_prob_bs)) :
tokenSoftmaxLlamaIO Logics B_Start_Loc B_Seqlen Prob_Out
stride_logic_h stride_logic_bs stride_prob_h stride_prob_bs BLOCK_SIZE
⊨ fun _ _ _ m₂ xs j => tokenSoftmaxRowSpec BLOCK_SIZE m₂ xs j
IO definition: tokenSoftmaxLlamaIOdef tokenSoftmaxLlamaIO
(Logics B_Start_Loc B_Seqlen Prob_Out : RegionName)
(stride_logic_h stride_logic_bs stride_prob_h stride_prob_bs
BLOCK_SIZE : Nat) :
MetaMasked2DKernelIO₁ where
kernel := token_softmax_surface Logics B_Start_Loc B_Seqlen Prob_Out
stride_logic_h stride_logic_bs stride_prob_h stride_prob_bs BLOCK_SIZE
mbuf1 := B_Start_Loc
mbuf2 := B_Seqlen
inp := Logics
out := Prob_Out
B := BLOCK_SIZE
mwin1 := fun pid₀ _ => pid₀
mwin2 := fun pid₀ _ => pid₀
read := fun _ pid₁ m₁ _ j =>
pid₁ * stride_logic_h + (m₁ + j.val) * stride_logic_bs
write := fun _ pid₁ m₁ _ j =>
pid₁ * stride_prob_h + (m₁ + j.val) * stride_prob_bs
mask := fun _ _ _ m₂ j => j.val < m₂
Referenced model/specification definitions: Python source |
|---|
triton_argmaxargmax_kernel_1_value_compute_correctMathematical execution | Configured model / stage Separate first-stage value/index, second-stage, and single-block dimension-reduction facts. First-stage IO fixes its flag; the host two-stage composition is not established by the bundle. | Statement and preconditionsspecification argmax_kernel_1_value_compute_correct
(inp mid_value : RegionName) (mid_index : Region .int)
(M BLOCK_SIZE : Nat)
(s : BlockState)
(hRegions : mid_value ≠ (Region.cast mid_index : RegionName)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := argmax_kernel_1 inp mid_value mid_index M BLOCK_SIZE Bool.false)
(initialState := s)
(write := fun _ : PUnit => some (mid_value, s.pid))
(expected := fun _ => argmaxKernel1ValueSpec s inp M BLOCK_SIZE)
Referenced model/specification definitions: Python source |
|---|
triton_argmaxargmax_kernel_1_index_compute_correctMathematical execution | Configured model / stage Separate first-stage value/index, second-stage, and single-block dimension-reduction facts. First-stage IO fixes its flag; the host two-stage composition is not established by the bundle. | Statement and preconditionsspecification argmax_kernel_1_index_compute_correct
(inp mid_value : RegionName) (mid_index : Region .int)
(M BLOCK_SIZE : Nat)
(s : BlockState) :
ComputeCorrect.Realizes_without_Rounding
(kernel := argmax_kernel_1 inp mid_value mid_index M BLOCK_SIZE Bool.false)
(initialState := s)
(write := fun _ : PUnit =>
some ((Region.cast mid_index : RegionName), s.pid))
(expected :=
fun _ : PUnit =>
(argmaxKernel1IndexSpec s inp M BLOCK_SIZE : Nat))
Referenced model/specification definitions: Python source |
|---|
triton_argmaxargmax_kernel_2_compute_correctMathematical execution | Configured model / stage Separate first-stage value/index, second-stage, and single-block dimension-reduction facts. First-stage IO fixes its flag; the host two-stage composition is not established by the bundle. | Statement and preconditionsspecification argmax_kernel_2_compute_correct
(mid_value mid_index out : RegionName)
(mid_size BLOCK_MID : Nat)
(s : BlockState) :
ComputeCorrect.Realizes_without_Rounding
(kernel := argmax_kernel_2 mid_value mid_index out mid_size BLOCK_MID)
(initialState := s)
(write := fun _ : PUnit => some (out, 0))
(expected := fun _ => argmaxKernel2Spec s mid_value mid_index mid_size BLOCK_MID)
Referenced model/specification definitions: Python source |
|---|
triton_argmaxargmax_kernel_dim_single_block_compute_correctMathematical execution | Configured model / stage Separate first-stage value/index, second-stage, and single-block dimension-reduction facts. First-stage IO fixes its flag; the host two-stage composition is not established by the bundle. | Statement and preconditionsspecification argmax_kernel_dim_single_block_compute_correct
(inp : RegionName) (out_index : Region .int)
(M N K BLOCK_M BLOCK_N : Nat)
(s : BlockState)
(hBN : 0 < BLOCK_N)
(hNpos : 0 < N)
(hNle : N ≤ BLOCK_N)
(hOutInj : Function.Injective
(fun i : Fin BLOCK_M => argmaxKernelOutOffset s K BLOCK_M i)) :
ComputeCorrect.Realizes_without_Rounding
(kernel := argmax_kernel inp out_index M N K BLOCK_M BLOCK_N Bool.false)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_M => s.pids 0 * BLOCK_M + i.val < M)
(fun i => ((Region.cast out_index : RegionName),
argmaxKernelOutOffset s K BLOCK_M i)))
(expected := fun i =>
argmaxKernelDimSingleBlockSpec s inp M N K BLOCK_M BLOCK_N i)
Referenced model/specification definitions: Python source |
|---|
triton_argmaxargmax_kernel_1_io_correctnessMathematical execution | Configured model / stage Separate first-stage value/index, second-stage, and single-block dimension-reduction facts. First-stage IO fixes its flag; the host two-stage composition is not established by the bundle. | Statement and preconditionsspecification argmax_kernel_1_io_correctness (inp mid_value : RegionName)
(mid_index : Region .int) (M BLOCK_SIZE : Nat) (hB : 0 < BLOCK_SIZE)
(hRegions : mid_value ≠ (Region.cast mid_index : RegionName)) :
ValueIndexTileKernelIO.Implements
(argmax1IO inp mid_value mid_index M BLOCK_SIZE)
(fun pid xs _ => argmaxValueSpecOf M BLOCK_SIZE pid xs)
(fun pid xs _ => argmaxIndexSpecOf M BLOCK_SIZE pid xs)
IO definition: argmax1IOdef argmax1IO (inp mid_value : RegionName) (mid_index : Region .int)
(M BLOCK_SIZE : Nat) : ValueIndexTileKernelIO where
kernel := argmax_kernel_1 inp mid_value mid_index M BLOCK_SIZE Bool.false
inp := inp
outVal := mid_value
outIdx := mid_index
shape := [BLOCK_SIZE]
read := fun pid idx => pid * BLOCK_SIZE + idx.1.val
writeVal := fun pid _ => pid
writeIdx := fun pid _ => pid
mask := fun pid idx => pid * BLOCK_SIZE + idx.1.val < M
writeMask := fun _pid idx => idx.1.val = 0
Referenced model/specification definitions: Python source |
|---|
triton_attentiontriton_attention_bwd_preprocess_genuine_output_summary_generalMathematical execution | Configured model / stage Separate backward-preprocess, forward, and backward-gradient value targets with explicit block/stride specializations. Forward IO adds frame only for its chosen surface; no whole training pipeline claim. | Statement and preconditionsspecification triton_attention_bwd_preprocess_genuine_output_summary_general
(Out DO L NewDO Delta : RegionName) (BLOCK_M D_HEAD : Nat) (s : BlockState)
(hND : NewDO ≠ Delta)
(hOutInj : Function.Injective
(fun idx : TileIndex [BLOCK_M, D_HEAD] =>
newdoOffset s BLOCK_M D_HEAD idx)) :
(∃ alg, (triton_attention_bwd_preprocess Out DO L NewDO Delta
BLOCK_M D_HEAD).toAlgorithm? = Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_attention_bwd_preprocess Out DO L NewDO Delta
BLOCK_M D_HEAD)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, D_HEAD] =>
some (NewDO, newdoOffset s BLOCK_M D_HEAD idx))
(expected := fun idx : TileIndex [BLOCK_M, D_HEAD] =>
bwdPreprocessNewDOSpecG s Out DO L BLOCK_M D_HEAD idx)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_attention_bwd_preprocess Out DO L NewDO Delta
BLOCK_M D_HEAD)
(initialState := s)
(write := fun i : Fin BLOCK_M => some (Delta, deltaOffset s BLOCK_M i))
(expected := fun i : Fin BLOCK_M =>
bwdPreprocessDeltaSpecG s Out DO L BLOCK_M D_HEAD i))
Referenced model/specification definitions: Python source |
|---|
triton_attentiontriton_attention_forward_output_summary_generalMathematical execution | Configured model / stage Separate backward-preprocess, forward, and backward-gradient value targets with explicit block/stride specializations. Forward IO adds frame only for its chosen surface; no whole training pipeline claim. | Statement and preconditionsspecification triton_attention_forward_output_summary_general
(Q K V L M Out : RegionName) (s : BlockState) (sc : ℝ)
(stride_qz stride_qh Z H N_CTX D0 BLOCK_M BLOCK_DMODEL BLOCK_N : Nat)
(hBM : 0 < BLOCK_M) (hBN : 0 < BLOCK_N) (hBD : 0 < BLOCK_DMODEL)
(hdvd : BLOCK_N ∣ (s.pids 0 + 1) * BLOCK_M)
(hbound : s.pids 1 * (stride_qh / BLOCK_DMODEL) + (s.pids 0 + 1) * BLOCK_M ≤ D0)
(hLOut : L ≠ Out) (hMOut : M ≠ Out) (hLM : M ≠ L)
(houtinj : Function.Injective (fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(s.pids 1 * (stride_qh / BLOCK_DMODEL) + s.pids 0 * BLOCK_M + idx.1.val) * BLOCK_DMODEL + idx.2.1.val * 1))
(hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (triton_attention_fwd_kernel Q K V L M Out sc
stride_qz stride_qh BLOCK_DMODEL 1 stride_qz stride_qh BLOCK_DMODEL 1
stride_qz stride_qh BLOCK_DMODEL 1 stride_qz stride_qh BLOCK_DMODEL 1
Z H N_CTX D0 BLOCK_M BLOCK_DMODEL BLOCK_N).toAlgorithm? = Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_attention_fwd_kernel Q K V L M Out sc
stride_qz stride_qh BLOCK_DMODEL 1 stride_qz stride_qh BLOCK_DMODEL 1
stride_qz stride_qh BLOCK_DMODEL 1 stride_qz stride_qh BLOCK_DMODEL 1
Z H N_CTX D0 BLOCK_M BLOCK_DMODEL BLOCK_N)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
active s (s.pids 1 * (stride_qh / BLOCK_DMODEL)) D0 BLOCK_M idx)
(fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
(Out, outOffset s (s.pids 1 * (stride_qh / BLOCK_DMODEL)) BLOCK_DMODEL 1 BLOCK_M idx)))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_DMODEL] =>
MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (fwdOutSpecG s Q K V (stride_qh / BLOCK_DMODEL) BLOCK_DMODEL
((s.pids 0 + 1) * BLOCK_M) BLOCK_M BLOCK_DMODEL sc idx))))) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_attention_fwd_kernel Q K V L M Out sc
stride_qz stride_qh BLOCK_DMODEL 1 stride_qz stride_qh BLOCK_DMODEL 1
stride_qz stride_qh BLOCK_DMODEL 1 stride_qz stride_qh BLOCK_DMODEL 1
Z H N_CTX D0 BLOCK_M BLOCK_DMODEL BLOCK_N)
(initialState := s)
(write := fun i : Fin BLOCK_M => some (L, lRowOffset s (s.pids 1) N_CTX BLOCK_M i))
(expected := fun i : Fin BLOCK_M =>
fwdLSpecG s Q K (stride_qh / BLOCK_DMODEL) BLOCK_DMODEL ((s.pids 0 + 1) * BLOCK_M) BLOCK_M BLOCK_DMODEL sc
(Nat.mul_pos (Nat.succ_pos _) hBM) i)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_attention_fwd_kernel Q K V L M Out sc
stride_qz stride_qh BLOCK_DMODEL 1 stride_qz stride_qh BLOCK_DMODEL 1
stride_qz stride_qh BLOCK_DMODEL 1 stride_qz stride_qh BLOCK_DMODEL 1
Z H N_CTX D0 BLOCK_M BLOCK_DMODEL BLOCK_N)
(initialState := s)
(write := fun i : Fin BLOCK_M => some (M, lRowOffset s (s.pids 1) N_CTX BLOCK_M i))
(expected := fun i : Fin BLOCK_M =>
fwdMSpecG s Q K (stride_qh / BLOCK_DMODEL) BLOCK_DMODEL ((s.pids 0 + 1) * BLOCK_M) BLOCK_M BLOCK_DMODEL sc
(Nat.mul_pos (Nat.succ_pos _) hBM) i))
Referenced model/specification definitions: Python source |
|---|
triton_attentiontriton_attention_bwd_grads_genuine_output_summary_generalMathematical execution | Configured model / stage Separate backward-preprocess, forward, and backward-gradient value targets with explicit block/stride specializations. Forward IO adds frame only for its chosen surface; no whole training pipeline claim. | Statement and preconditionsspecification triton_attention_bwd_grads_genuine_output_summary_general (H stride_qz stride_qh : Nat) (base : Nat)
(Q K V Out DO DQ DK DV L M Delta : RegionName) (s : BlockState) (sc : ℝ)
(BM BD D0 nb : Nat)
-- slots the kernel ignores (`_Z`, `_stride_k{z,h}`, `_stride_v{z,h}`,
-- `_BLOCK_N`): universally quantified, since nothing depends on them
(Z skz skh svz svh BN : Nat)
(hBM : 0 < BM) (hBD : 0 < BD) (hnb : 0 < nb) (hbdvd : BD ∣ base)
(hbound : base / BD + nb * BM ≤ D0)
(hbase : (s.pids 0 / H) * (stride_qz / BD) + (s.pids 0 % H) * (stride_qh / BD) = base / BD)
(hQDQ : Q ≠ DQ) (hKDQ : K ≠ DQ) (hVDQ : V ≠ DQ) (hDODQ : DO ≠ DQ)
(hMDQ : M ≠ DQ) (hDeDQ : Delta ≠ DQ)
(hDVDQ : DV ≠ DQ) (hDKDQ : DK ≠ DQ) (hDVDK : DV ≠ DK) (hDKDV : DK ≠ DV)
(hin : ∀ R : RegionName, R = Q ∨ R = K ∨ R = V ∨ R = DO ∨ R = M ∨ R = Delta →
R ≠ DV ∧ R ≠ DK ∧ R ≠ DQ)
(hundef : ∀ rg o, s.undef rg o = 0) :
(∃ alg, (triton_attention_bwd_kernel Q K V Out DO DQ DK DV L M Delta sc
stride_qz stride_qh BD 1 skz skh BD 1 svz svh BD 1
Z H (BM * nb) D0 nb BM BD BN).toAlgorithm? = Except.ok alg) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_attention_bwd_kernel Q K V Out DO DQ DK DV L M Delta sc
stride_qz stride_qh BD 1 skz skh BD 1 svz svh BD 1
Z H (BM * nb) D0 nb BM BD BN)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BM * nb, BD] => idx.1.val < nb * BM)
(fun idx : TileIndex [BM * nb, BD] => (DQ, base + idx.1.val * BD + idx.2.1.val)))
(expected := fun idx : TileIndex [BM * nb, BD] =>
bwdKernelDQSpecG base s Q K V DO M Delta DQ BD (BM * nb) sc idx.1.val idx.2.1.val)) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_attention_bwd_kernel Q K V Out DO DQ DK DV L M Delta sc
stride_qz stride_qh BD 1 skz skh BD 1 svz svh BD 1
Z H (BM * nb) D0 nb BM BD BN)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BM * nb, BD] => idx.1.val < nb * BM)
(fun idx : TileIndex [BM * nb, BD] => (DV, base + idx.1.val * BD + idx.2.1.val)))
(expected := fun idx : TileIndex [BM * nb, BD] =>
MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (∑ I : Fin (BM * nb),
bwdFp16 (bwdKernelPG base s Q K M BD (BM * nb) sc I.val idx.1.val) *
bwdKernelDOG base s DO BD I.val idx.2.1.val))))) ∧
(ComputeCorrect.Realizes_without_Rounding
(kernel := triton_attention_bwd_kernel Q K V Out DO DQ DK DV L M Delta sc
stride_qz stride_qh BD 1 skz skh BD 1 svz svh BD 1
Z H (BM * nb) D0 nb BM BD BN)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun idx : TileIndex [BM * nb, BD] => idx.1.val < nb * BM)
(fun idx : TileIndex [BM * nb, BD] => (DK, base + idx.1.val * BD + idx.2.1.val)))
(expected := fun idx : TileIndex [BM * nb, BD] =>
MemCell.of .fp16 (FloatDType.real.cast FloatDType.fp16
(some (∑ I : Fin (BM * nb),
bwdFp16 (bwdKernelDSG base s Q K V DO M Delta BD (BM * nb) sc I.val idx.1.val) *
bwdKernelQG base s Q BD I.val idx.2.1.val)))))
Referenced model/specification definitions: Python source |
|---|
triton_attentiontriton_attention_forward_io_correctnessAbstract cast/store rounding | Configured model / stage Separate backward-preprocess, forward, and backward-gradient value targets with explicit block/stride specializations. Forward IO adds frame only for its chosen surface; no whole training pipeline claim. | Statement and preconditionsspecification triton_attention_forward_io_correctness (R : RoundingModel)
(hfp16 : R.round .fp16 = id)
(Q K V L M Out : RegionName) (sc : ℝ)
(stride_qz stride_qh Z H N_CTX D0 BLOCK_M BLOCK_DMODEL BLOCK_N : Nat)
(hBM : 0 < BLOCK_M) (hBN : 0 < BLOCK_N) (hBD : 0 < BLOCK_DMODEL)
(hBNM : BLOCK_N ∣ BLOCK_M)
(hLOut : L ≠ Out) (hMOut : M ≠ Out) (hLM : M ≠ L) :
tritonAttentionFwdIO Q K V L M Out sc
stride_qz stride_qh Z H N_CTX D0 BLOCK_M BLOCK_DMODEL BLOCK_N ⊨[R]
fun p₀ p₁ _ xs ys zs =>
(fun j : Fin (BLOCK_M * BLOCK_DMODEL) =>
attentionRealCausalBlock (p₀ * BLOCK_M)
(taIOqT BLOCK_M BLOCK_DMODEL (D0 / BLOCK_N) xs)
(taIOkvT BLOCK_N BLOCK_DMODEL (D0 / BLOCK_N) ((p₀ + 1) * BLOCK_M) ys)
(taIOkvT BLOCK_N BLOCK_DMODEL (D0 / BLOCK_N) ((p₀ + 1) * BLOCK_M) zs)
sc (Lane2D.decode j),
fun i : Fin BLOCK_M =>
taIOLSpecT p₀ ((p₀ + 1) * BLOCK_M) BLOCK_M BLOCK_DMODEL sc
(taIOqT BLOCK_M BLOCK_DMODEL (D0 / BLOCK_N) xs)
(taIOkvT BLOCK_N BLOCK_DMODEL (D0 / BLOCK_N) ((p₀ + 1) * BLOCK_M) ys) i,
fun i : Fin BLOCK_M =>
taIOMSpecT p₀ ((p₀ + 1) * BLOCK_M) BLOCK_M BLOCK_DMODEL sc
(taIOqT BLOCK_M BLOCK_DMODEL (D0 / BLOCK_N) xs)
(taIOkvT BLOCK_N BLOCK_DMODEL (D0 / BLOCK_N) ((p₀ + 1) * BLOCK_M) ys) i)
IO definition: tritonAttentionFwdIOdef tritonAttentionFwdIO (Q K V L M Out : RegionName) (sc : ℝ)
(stride_qz stride_qh Z H N_CTX D0 BLOCK_M BLOCK_DMODEL BLOCK_N : Nat) :
StreamMasked3DKernelIO₃ₓ₃ where
kernel := triton_attention_fwd_kernel Q K V L M Out sc
stride_qz stride_qh BLOCK_DMODEL 1 stride_qz stride_qh BLOCK_DMODEL 1
stride_qz stride_qh BLOCK_DMODEL 1 stride_qz stride_qh BLOCK_DMODEL 1
Z H N_CTX D0 BLOCK_M BLOCK_DMODEL BLOCK_N
inp1 := Q
inp2 := K
inp3 := V
out1 := Out
out2 := L
out3 := M
T := D0 / BLOCK_N
B1 := BLOCK_M * BLOCK_DMODEL
B2 := BLOCK_N * BLOCK_DMODEL
B3 := BLOCK_N * BLOCK_DMODEL
C1 := BLOCK_M * BLOCK_DMODEL
C2 := BLOCK_M
C3 := BLOCK_M
out1DType := .fp16
read1 := fun p₀ p₁ _ _ j =>
(p₁ * (stride_qh / BLOCK_DMODEL) + p₀ * BLOCK_M + j.val / BLOCK_DMODEL) * BLOCK_DMODEL
+ j.val % BLOCK_DMODEL
read2 := fun _ p₁ _ t j =>
(p₁ * (stride_qh / BLOCK_DMODEL) + t.val * BLOCK_N + j.val / BLOCK_DMODEL) * BLOCK_DMODEL
+ j.val % BLOCK_DMODEL
read3 := fun _ p₁ _ t j =>
(p₁ * (stride_qh / BLOCK_DMODEL) + t.val * BLOCK_N + j.val / BLOCK_DMODEL) * BLOCK_DMODEL
+ j.val % BLOCK_DMODEL
write1 := fun p₀ p₁ _ j =>
(p₁ * (stride_qh / BLOCK_DMODEL) + p₀ * BLOCK_M + j.val / BLOCK_DMODEL) * BLOCK_DMODEL
+ (j.val % BLOCK_DMODEL) * 1
write2 := fun p₀ p₁ _ i => p₁ * N_CTX + (p₀ * BLOCK_M + i.val)
write3 := fun p₀ p₁ _ i => p₁ * N_CTX + (p₀ * BLOCK_M + i.val)
mask1 := fun _ _ _ _ _ => True
mask2 := fun p₀ _ _ t _ => t.val * BLOCK_N < (p₀ + 1) * BLOCK_M
mask3 := fun p₀ _ _ t _ => t.val * BLOCK_N < (p₀ + 1) * BLOCK_M
writeMask1 := fun _ _ _ _ => True
writeMask2 := fun _ _ _ _ => True
writeMask3 := fun _ _ _ _ => True
pre := fun p₀ p₁ _ =>
p₁ * (stride_qh / BLOCK_DMODEL) + (p₀ + 1) * BLOCK_M ≤ D0
Referenced model/specification definitions: Python source |
|---|
triton_conv2d_fwdconv2d_output_summaryMathematical execution | Configured model / stage Configured forward convolution with the IO binding fixing the shown branch flag and in-channel/block relation. Groups, padding, strides, and tf32 have only the stated mathematical semantics. | Statement and preconditionsspecification conv2d_output_summary
(Input Weight Output : RegionName)
(batch_dim in_feat_dim in_height in_width out_feat_dim out_height out_width
IBS IIFS IHS IWS WOFS WIFS WHS WWS OBS OOFS OHS OWS
KH KW SH SW PH PW groups : Nat) (tf32 : Bool) (BHW BIN OF numCBlocks : Nat)
(s : BlockState) (hBIN : 0 < BIN) (hundef : ∀ rg o, s.undef rg o = 0)
(hIGD : in_feat_dim / groups = BIN * numCBlocks)
(hOutInj : Function.Injective (outputOffset s BHW OF out_height out_width OBS OOFS OHS OWS (out_feat_dim / groups))) :
(∃ alg, (conv2d_forward_surface Input Weight Output batch_dim in_feat_dim in_height in_width
out_feat_dim out_height out_width IBS IIFS IHS IWS WOFS WIFS WHS WWS OBS OOFS OHS OWS
KH KW SH SW PH PW groups Bool.true tf32 BHW BIN OF).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := conv2d_forward_surface Input Weight Output batch_dim in_feat_dim in_height in_width
out_feat_dim out_height out_width IBS IIFS IHS IWS WOFS WIFS WHS WWS OBS OOFS OHS OWS
KH KW SH SW PH PW groups Bool.true tf32 BHW BIN OF)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s BHW OF batch_dim out_height out_width (out_feat_dim / groups))
(fun idx => (Output, outputOffset s BHW OF out_height out_width OBS OOFS OHS OWS (out_feat_dim / groups) idx)))
(expected := fun idx : TileIndex [BHW, OF] =>
convSpec s Input Weight BHW BIN OF batch_dim in_height in_width
IBS IIFS IHS IWS WOFS WIFS WHS WWS out_height out_width (BIN * numCBlocks) (out_feat_dim / groups)
SH SW PH PW KH KW numCBlocks idx.1 idx.2.1)
Referenced model/specification definitions: Python source |
|---|
triton_conv2d_fwdtriton_conv2d_fwd_io_correctnessAbstract cast/store rounding | Configured model / stage Configured forward convolution with the IO binding fixing the shown branch flag and in-channel/block relation. Groups, padding, strides, and tf32 have only the stated mathematical semantics. | Statement and preconditionsspecification triton_conv2d_fwd_io_correctness (R : RoundingModel)
(Input Weight Output : RegionName)
(batch_dim in_feat_dim in_height in_width out_feat_dim out_height out_width
IBS IIFS IHS IWS WOFS WIFS WHS WWS OBS OOFS OHS OWS
KH KW SH SW PH PW groups : Nat) (tf32 : Bool) (BHW BIN OF numCBlocks : Nat)
(hfp16 : R.round .fp16 = id) (hBIN : 0 < BIN)
(hIGD : in_feat_dim / groups = BIN * numCBlocks)
(hOutInj : ∀ pid₀ pid₁ pid₂ : Nat,
Function.Injective (fun idx : TileIndex [BHW, OF] =>
pOutAddr pid₀ pid₁ pid₂ BHW OF out_height out_width OBS OOFS OHS OWS
(out_feat_dim / groups) idx.1.val idx.2.1.val)) :
triton_conv2d_fwd_IO Input Weight Output batch_dim in_feat_dim in_height in_width
out_feat_dim out_height out_width IBS IIFS IHS IWS WOFS WIFS WHS WWS
OBS OOFS OHS OWS KH KW SH SW PH PW groups tf32 BHW BIN OF numCBlocks ⊨[R]
fun pid₀ pid₁ _ _ xs ys l =>
convStreamSum pid₀ pid₁ batch_dim in_height in_width out_height out_width
(out_feat_dim / groups) SH SW PH PW KH KW BHW BIN OF numCBlocks xs ys l
IO definition: triton_conv2d_fwd_IOdef triton_conv2d_fwd_IO (Input Weight Output : RegionName)
(batch_dim in_feat_dim in_height in_width out_feat_dim out_height out_width
IBS IIFS IHS IWS WOFS WIFS WHS WWS OBS OOFS OHS OWS
KH KW SH SW PH PW groups : Nat) (tf32 : Bool) (BHW BIN OF numCBlocks : Nat) :
StreamMetaMasked3DKernelIO₂ where
kernel := conv2d_forward_surface Input Weight Output batch_dim in_feat_dim in_height in_width
out_feat_dim out_height out_width IBS IIFS IHS IWS WOFS WIFS WHS WWS OBS OOFS OHS OWS
KH KW SH SW PH PW groups Bool.true tf32 BHW BIN OF
inp1 := Input
inp2 := Weight
out := Output
nMeta := 0
sty := Fin.elim0
mbuf := Fin.elim0
mwin := Fin.elim0
T := KH * KW * numCBlocks
B1 := BHW * BIN
B2 := BIN * OF
C := BHW * OF
outDType := .real
read1 := fun pid₀ _ pid₂ _ t j =>
pInAddr pid₀ pid₂ IBS IIFS IHS IWS out_height out_width BHW (BIN * numCBlocks) SH SW PH PW
(convStepH KW numCBlocks t.val) (convStepW KW numCBlocks t.val)
(convStepCB numCBlocks t.val * BIN) (j.val / BIN) (j.val % BIN)
read2 := fun _ pid₁ pid₂ _ t j =>
pWAddr pid₁ pid₂ WOFS WIFS WHS WWS OF (out_feat_dim / groups)
(convStepH KW numCBlocks t.val) (convStepW KW numCBlocks t.val)
(convStepCB numCBlocks t.val * BIN) (j.val / OF) (j.val % OF)
write := fun pid₀ pid₁ pid₂ _ j =>
pOutAddr pid₀ pid₁ pid₂ BHW OF out_height out_width OBS OOFS OHS OWS
(out_feat_dim / groups) (j.val / OF) (j.val % OF)
mask1 := fun pid₀ _ _ _ t j =>
pInMask pid₀ batch_dim in_height in_width out_height out_width BHW (BIN * numCBlocks)
SH SW PH PW (convStepH KW numCBlocks t.val) (convStepW KW numCBlocks t.val)
(convStepCB numCBlocks t.val * BIN) (j.val / BIN) (j.val % BIN)
mask2 := fun _ pid₁ _ _ t j =>
pWMask pid₁ (BIN * numCBlocks) (out_feat_dim / groups) OF
(convStepCB numCBlocks t.val * BIN) (j.val / OF) (j.val % OF)
writeMask := fun pid₀ pid₁ _ _ j =>
pActive pid₀ pid₁ BHW OF batch_dim out_height out_width (out_feat_dim / groups)
(j.val / OF) (j.val % OF)
Referenced model/specification definitions: Python source |
|---|
triton_linear_activationtriton_linear_activation_output_summary_generalMathematical execution | Configured model / stage Configured linear/bias/activation and optional activation-input-save paths. The named activation functions and reduction/layout conditions determine the claim; no host autotuning or approximation-error certificate. | Statement and preconditionsspecification triton_linear_activation_output_summary_general
(C ACT_INPUTS A B bias : RegionName) (s : BlockState)
(M N output_m_stride output_n_stride act_inputs_m_stride act_inputs_n_stride
a_m_stride a_k_stride b_n_stride b_k_stride
BLOCK_M GROUP_M BLOCK_N BLOCK_K numKBlocks : Nat)
(HAS_BIAS SHOULD_SAVE_ACT_INPUTS : Bool) (ACTIVATION : String)
(hFitM : blockMIdx (s.pids 0) M N BLOCK_M BLOCK_N GROUP_M * BLOCK_M + BLOCK_M ≤ M)
(hFitN : blockNIdx (s.pids 0) M N BLOCK_M BLOCK_N GROUP_M * BLOCK_N + BLOCK_N ≤ N)
(hsno : output_n_stride = 1) (hble : BLOCK_N ≤ output_m_stride) :
-- (1) the surface lowers to the algorithm layer
(∃ alg, (kernel_fma_surface C ACT_INPUTS A B bias M N output_m_stride output_n_stride
act_inputs_m_stride act_inputs_n_stride a_m_stride a_k_stride b_n_stride b_k_stride
BLOCK_M GROUP_M BLOCK_N BLOCK_K numKBlocks
HAS_BIAS SHOULD_SAVE_ACT_INPUTS ACTIVATION).toAlgorithm? = Except.ok alg) ∧
-- (2) C: genuine fused linear + activation
ComputeCorrect.Realizes_without_Rounding
(kernel := kernel_fma_surface C ACT_INPUTS A B bias M N output_m_stride output_n_stride
act_inputs_m_stride act_inputs_n_stride a_m_stride a_k_stride b_n_stride b_k_stride
BLOCK_M GROUP_M BLOCK_N BLOCK_K numKBlocks
HAS_BIAS SHOULD_SAVE_ACT_INPUTS ACTIVATION)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
some (C, cOffset s M N BLOCK_M BLOCK_N GROUP_M output_m_stride output_n_stride idx))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
applyActivation ACTIVATION
(linearSpec s A B bias M N BLOCK_M BLOCK_N GROUP_M a_m_stride a_k_stride
b_k_stride b_n_stride BLOCK_K numKBlocks HAS_BIAS idx.1 idx.2.1)) ∧
-- (3) ACT_INPUTS (when saving): the genuine un-activated pre-activation values
(SHOULD_SAVE_ACT_INPUTS = Bool.true → ACT_INPUTS ≠ C →
act_inputs_n_stride = 1 → BLOCK_N ≤ act_inputs_m_stride →
ComputeCorrect.Realizes_without_Rounding
(kernel := kernel_fma_surface C ACT_INPUTS A B bias M N output_m_stride output_n_stride
act_inputs_m_stride act_inputs_n_stride a_m_stride a_k_stride b_n_stride b_k_stride
BLOCK_M GROUP_M BLOCK_N BLOCK_K numKBlocks
HAS_BIAS SHOULD_SAVE_ACT_INPUTS ACTIVATION)
(initialState := s)
(write := fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
some (ACT_INPUTS, actOffset s M N BLOCK_M BLOCK_N GROUP_M
act_inputs_m_stride act_inputs_n_stride idx))
(expected := fun idx : TileIndex [BLOCK_M, BLOCK_N] =>
linearSpec s A B bias M N BLOCK_M BLOCK_N GROUP_M a_m_stride a_k_stride
b_k_stride b_n_stride BLOCK_K numKBlocks HAS_BIAS idx.1 idx.2.1))
Referenced model/specification definitions: Python source |
|---|
triton_linear_activationtriton_linear_activation_io_correctnessAbstract cast/store rounding | Configured model / stage Configured linear/bias/activation and optional activation-input-save paths. The named activation functions and reduction/layout conditions determine the claim; no host autotuning or approximation-error certificate. | Statement and preconditionsspecification triton_linear_activation_io_correctness (R : RoundingModel)
(C ACT_INPUTS A B bias : RegionName)
(M N output_m_stride output_n_stride act_inputs_m_stride act_inputs_n_stride
a_m_stride a_k_stride b_n_stride b_k_stride
BLOCK_M GROUP_M BLOCK_N BLOCK_K numKBlocks : Nat)
(HAS_BIAS SHOULD_SAVE_ACT_INPUTS : Bool) (ACTIVATION : String)
(hT : 0 < numKBlocks)
(hsno : output_n_stride = 1) (hble : BLOCK_N ≤ output_m_stride)
(hsain : act_inputs_n_stride = 1) (hale : BLOCK_N ≤ act_inputs_m_stride)
(hne : ACT_INPUTS ≠ C) :
tlaIO C ACT_INPUTS A B bias M N output_m_stride output_n_stride
act_inputs_m_stride act_inputs_n_stride a_m_stride a_k_stride b_n_stride b_k_stride
BLOCK_M GROUP_M BLOCK_N BLOCK_K numKBlocks
HAS_BIAS SHOULD_SAVE_ACT_INPUTS ACTIVATION ⊨[R]
fun p₀ _ _ xs ys zs =>
(fun l => applyActivation ACTIVATION
(tlaLinIO p₀ M N BLOCK_M BLOCK_N GROUP_M BLOCK_K numKBlocks HAS_BIAS xs ys zs l),
fun l => tlaLinIO p₀ M N BLOCK_M BLOCK_N GROUP_M BLOCK_K numKBlocks HAS_BIAS xs ys zs l)
IO definition: tlaIOnoncomputable def tlaIO (C ACT_INPUTS A B biasR : RegionName)
(M N smo sno saim sain sam sak sbn sbk BM GM BN BK numKBlocks : Nat)
(HAS_BIAS SHOULD_SAVE_ACT_INPUTS : Bool) (ACTIVATION : String) :
StreamMasked3DKernelIO₃ₓ₂ where
kernel := kernel_fma_surface C ACT_INPUTS A B biasR M N smo sno saim sain
sam sak sbn sbk BM GM BN BK numKBlocks HAS_BIAS SHOULD_SAVE_ACT_INPUTS ACTIVATION
inp1 := A
inp2 := B
inp3 := biasR
out1 := C
out2 := ACT_INPUTS
T := numKBlocks
B1 := BM * BK
B2 := BK * BN
B3 := BN
C1 := BM * BN
C2 := BM * BN
out1DType := .real
out2DType := .real
pre := fun p₀ _ _ =>
blockMIdx p₀ M N BM BN GM * BM + BM ≤ M ∧ blockNIdx p₀ M N BM BN GM * BN + BN ≤ N
read1 := fun p₀ _ _ t l => tlaRowIdx p₀ M N BM BN GM (l.val / BK) * sam + l.val % BK * sak + t.val * BK * sak
read2 := fun p₀ _ _ t l => l.val / BN * sbk + tlaColIdx p₀ M N BM BN GM (l.val % BN) * sbn + t.val * BK * sbk
read3 := fun p₀ _ _ _ j => tlaColIdx p₀ M N BM BN GM j.val
write1 := fun p₀ _ _ l => tlaRowIdx p₀ M N BM BN GM (l.val / BN) * smo + tlaColIdx p₀ M N BM BN GM (l.val % BN) * sno
write2 := fun p₀ _ _ l => tlaRowIdx p₀ M N BM BN GM (l.val / BN) * saim + tlaColIdx p₀ M N BM BN GM (l.val % BN) * sain
mask1 := fun _ _ _ _ _ => True
mask2 := fun _ _ _ _ _ => True
mask3 := fun p₀ _ _ _ j => HAS_BIAS = Bool.true ∧ tlaColIdx p₀ M N BM BN GM j.val < N
writeMask1 := fun _ _ _ _ => True
writeMask2 := fun _ _ _ _ => SHOULD_SAVE_ACT_INPUTS = Bool.true
Referenced model/specification definitions: Python source |
|---|
triton_matmultriton_matmul_f16_closed_form_correctMathematical execution | Configured model / stage Separate fp16/f8e4 output models with whole reduction blocks and contiguous output-column layout. Abstract dtype rounding is not concrete hardware fp8/IEEE behavior. | Statement and preconditionsspecification triton_matmul_f16_closed_form_correct
(A B C : RegionName) (s : BlockState)
(M N BM BN GM sam sak sbk sbn scm scn BLOCK_K numKBlocks : Nat) (K : Nat)
(hK : K = BLOCK_K * numKBlocks)
(hcn : scn = 1) (hbnle : BN ≤ scm)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := triton_matmul_f16_surface A B C M N K sam sak sbk sbn scm scn
BM BN BLOCK_K GM numKBlocks)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BM BN GM)
(fun idx => (C, cOffset s M N BM BN GM scm scn idx)))
(expected := fun idx : TileIndex [BM, BN] =>
MemCell.of .fp16
(FloatDType.real.cast FloatDType.fp16
(some (matmulSpec s A B M N BM BN GM sam sak sbk sbn BLOCK_K numKBlocks idx.1 idx.2.1))))
Referenced model/specification definitions: Python source |
|---|
triton_matmultriton_matmul_f8_closed_form_correctMathematical execution | Configured model / stage Separate fp16/f8e4 output models with whole reduction blocks and contiguous output-column layout. Abstract dtype rounding is not concrete hardware fp8/IEEE behavior. | Statement and preconditionsspecification triton_matmul_f8_closed_form_correct
(A B C : RegionName) (s : BlockState)
(M N BM BN GM sam sak sbk sbn scm scn BLOCK_K numKBlocks : Nat) (K : Nat)
(hK : K = BLOCK_K * numKBlocks)
(hcn : scn = 1) (hbnle : BN ≤ scm)
(hundef : ∀ rg o, s.undef rg o = 0) :
ComputeCorrect.Realizes_without_Rounding
(kernel := triton_matmul_f8_surface A B C M N K sam sak sbk sbn scm scn
BM BN BLOCK_K GM numKBlocks)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(active s M N BM BN GM)
(fun idx => (C, cOffset s M N BM BN GM scm scn idx)))
(expected := fun idx : TileIndex [BM, BN] =>
MemCell.of .f8e4
(FloatDType.real.cast FloatDType.f8e4
(some (matmulSpec s A B M N BM BN GM sam sak sbk sbn BLOCK_K numKBlocks idx.1 idx.2.1))))
Referenced model/specification definitions: Python source |
|---|
triton_matmultriton_matmul_f16_io_correctnessAbstract cast/store rounding | Configured model / stage Separate fp16/f8e4 output models with whole reduction blocks and contiguous output-column layout. Abstract dtype rounding is not concrete hardware fp8/IEEE behavior. | Statement and preconditionsspecification triton_matmul_f16_io_correctness (R : RoundingModel)
(A B C : RegionName)
(M N K sam sak sbk sbn scm scn BM BN BK GM numKBlocks : Nat)
(hK : K = BK * numKBlocks) (hcn : scn = 1) (hBN : BN ≤ scm) :
tritonMatmulF16IO A B C M N K sam sak sbk sbn scm scn BM BN BK GM numKBlocks
⊨[R] fun _ _ xs ys l =>
∑ t : Fin numKBlocks, ∑ e : Fin BK,
xs t (aLane BM BN BK l e) * ys t (bLane BM BN BK l e)
IO definition: tritonMatmulF16IOdef tritonMatmulF16IO (A B C : RegionName)
(M N K sam sak sbk sbn scm scn BM BN BK GM numKBlocks : Nat) :
StreamMasked2DKernelIO₂ where
kernel := triton_matmul_f16_surface A B C M N K sam sak sbk sbn scm scn BM BN BK GM
numKBlocks
inp1 := A
inp2 := B
out := C
T := numKBlocks
B1 := BM * BK
B2 := BK * BN
C := BM * BN
outDType := .fp16
read1 := fun p₀ _ t l =>
clampIdx (pidM p₀ M N BM BN GM * BM + l.val / BK) M * sam + (t.val * BK + l.val % BK) * sak
read2 := fun p₀ _ t l =>
(t.val * BK + l.val / BN) * sbk + clampIdx (pidN p₀ M N BM BN GM * BN + l.val % BN) N * sbn
write := fun p₀ _ l =>
scm * (pidM p₀ M N BM BN GM * BM + l.val / BN) + scn * (pidN p₀ M N BM BN GM * BN + l.val % BN)
mask1 := fun _ _ t l => t.val * BK + l.val % BK < K
mask2 := fun _ _ t l => t.val * BK + l.val / BN < K
writeMask := fun p₀ _ l =>
pidM p₀ M N BM BN GM * BM + l.val / BN < M ∧ pidN p₀ M N BM BN GM * BN + l.val % BN < N
Referenced model/specification definitions: Python source |
|---|
triton_matmultriton_matmul_f8_io_correctnessAbstract cast/store rounding | Configured model / stage Separate fp16/f8e4 output models with whole reduction blocks and contiguous output-column layout. Abstract dtype rounding is not concrete hardware fp8/IEEE behavior. | Statement and preconditionsspecification triton_matmul_f8_io_correctness (R : RoundingModel)
(A B C : RegionName)
(M N K sam sak sbk sbn scm scn BM BN BK GM numKBlocks : Nat)
(hK : K = BK * numKBlocks) (hcn : scn = 1) (hBN : BN ≤ scm) :
tritonMatmulF8IO A B C M N K sam sak sbk sbn scm scn BM BN BK GM numKBlocks
⊨[R] fun _ _ xs ys l =>
∑ t : Fin numKBlocks, ∑ e : Fin BK,
xs t (aLane BM BN BK l e) * ys t (bLane BM BN BK l e)
IO definition: tritonMatmulF8IOdef tritonMatmulF8IO (A B C : RegionName)
(M N K sam sak sbk sbn scm scn BM BN BK GM numKBlocks : Nat) :
StreamMasked2DKernelIO₂ where
kernel := triton_matmul_f8_surface A B C M N K sam sak sbk sbn scm scn BM BN BK GM
numKBlocks
inp1 := A
inp2 := B
out := C
T := numKBlocks
B1 := BM * BK
B2 := BK * BN
C := BM * BN
outDType := .f8e4
read1 := fun p₀ _ t l =>
clampIdx (pidM p₀ M N BM BN GM * BM + l.val / BK) M * sam + (t.val * BK + l.val % BK) * sak
read2 := fun p₀ _ t l =>
(t.val * BK + l.val / BN) * sbk + clampIdx (pidN p₀ M N BM BN GM * BN + l.val % BN) N * sbn
write := fun p₀ _ l =>
scm * (pidM p₀ M N BM BN GM * BM + l.val / BN) + scn * (pidN p₀ M N BM BN GM * BN + l.val % BN)
mask1 := fun _ _ t l => t.val * BK + l.val % BK < K
mask2 := fun _ _ t l => t.val * BK + l.val / BN < K
writeMask := fun p₀ _ l =>
pidM p₀ M N BM BN GM * BM + l.val / BN < M ∧ pidN p₀ M N BM BN GM * BN + l.val % BN < N
Referenced model/specification definitions: Python source |
|---|
triton_mul2mul2_kernel_correctnessMathematical execution | Configured model / stage Separate masked out-of-place and in-place doubling contracts. The in-place theorem relates the final values to the initially loaded window. | Statement and preconditionsspecification mul2_kernel_correctness
(in_ptr0 out_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) :
mul2IO in_ptr0 out_ptr n_elements BLOCK_SIZE
⊨ fun xs i => 2 * xs i
IO definition: mul2IOdef mul2IO (in_ptr0 out_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) : MaskedKernelIO₁ where
kernel := mul2_kernel in_ptr0 out_ptr n_elements BLOCK_SIZE
inp := in_ptr0
out := out_ptr
B := BLOCK_SIZE
read := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < n_elements
Referenced model/specification definitions: Python source |
|---|
triton_mul2mul2_inplace_kernel_correctnessMathematical execution | Configured model / stage Separate masked out-of-place and in-place doubling contracts. The in-place theorem relates the final values to the initially loaded window. | Statement and preconditionsspecification mul2_inplace_kernel_correctness
(ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) :
mul2InplaceIO ptr n_elements BLOCK_SIZE
⊨ fun xs i => 2 * xs i
IO definition: mul2InplaceIOdef mul2InplaceIO (ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) : MaskedKernelIO₁ where
kernel := mul2_inplace_kernel ptr n_elements BLOCK_SIZE
inp := ptr
out := ptr
B := BLOCK_SIZE
read := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < n_elements
Referenced model/specification definitions: Python source |
|---|
triton_softmaxsoftmax_kernel_correctnessMathematical execution | Configured model / stage One masked row softmax with positive block size and explicit row strides, under mathematical exp/division semantics. | Statement and preconditionsspecification softmax_kernel_correctness
(output_ptr input_ptr : RegionName)
(input_row_stride output_row_stride n_cols BLOCK_SIZE : Nat)
(hB : 0 < BLOCK_SIZE) :
softmaxIO output_ptr input_ptr input_row_stride output_row_stride
n_cols BLOCK_SIZE ⊨
fun xs i => softmaxSpec n_cols BLOCK_SIZE xs i
IO definition: softmaxIOdef softmaxIO (output_ptr input_ptr : RegionName)
(input_row_stride output_row_stride n_cols BLOCK_SIZE : Nat) :
MaskedKernelIO₁ where
kernel := softmax_kernel output_ptr input_ptr input_row_stride
output_row_stride n_cols BLOCK_SIZE
inp := input_ptr
out := output_ptr
B := BLOCK_SIZE
read := fun pid => pid * input_row_stride
write := fun pid => pid * output_row_stride
mask := fun _ j => j.val < n_cols
Referenced model/specification definitions: Python source |
|---|
var_len_copyvar_len_copy_kernel_triton_small_length_output_summaryMathematical execution | Configured model / stage The full-surface summary is restricted to small lengths; the IO theorem copies one explicitly indexed chunk. A general multi-chunk host copy is not inferred. | Statement and preconditionsspecification var_len_copy_kernel_triton_small_length_output_summary
(old_a_start old_a_len : Region .nat) (old_a_location : RegionName)
(new_a_start : Region .nat) (new_a_location : RegionName)
(BLOCK_SIZE : Nat)
(s : BlockState)
(hBS : 0 < BLOCK_SIZE)
(hLen :
s.readMemValue .nat (Region.cast old_a_len : RegionName) (s.pids 0)
≤ BLOCK_SIZE)
(hLenPos :
0 < s.readMemValue .nat (Region.cast old_a_len : RegionName) (s.pids 0))
(hOutInj : Function.Injective
(fun i : Fin BLOCK_SIZE =>
s.readMemValue .nat (Region.cast new_a_start : RegionName) (s.pids 0)
+ i.val)) :
(∃ alg, (var_len_copy_kernel_triton old_a_start old_a_len old_a_location
new_a_start new_a_location BLOCK_SIZE).toAlgorithm? = Except.ok alg) ∧
ComputeCorrect.Realizes_without_Rounding
(kernel := var_len_copy_kernel_triton old_a_start old_a_len old_a_location
new_a_start new_a_location BLOCK_SIZE)
(initialState := s)
(write := ComputeCorrect.WriteMap.writeIf
(fun i : Fin BLOCK_SIZE =>
i.val < s.readMemValue .nat (Region.cast old_a_len : RegionName) (s.pids 0))
(fun i =>
(new_a_location,
s.readMemValue .nat (Region.cast new_a_start : RegionName) (s.pids 0)
+ i.val)))
(expected := fun i =>
s.readMem old_a_location
(s.readMemValue .nat (Region.cast old_a_start : RegionName) (s.pids 0)
+ i.val))
Referenced model/specification definitions: Python source |
|---|
var_len_copyvar_len_copy_one_chunk_io_correctnessMathematical execution | Configured model / stage The full-surface summary is restricted to small lengths; the IO theorem copies one explicitly indexed chunk. A general multi-chunk host copy is not inferred. | Statement and preconditionsspecification var_len_copy_one_chunk_io_correctness
(old_a_start old_a_len old_a_location new_a_start new_a_location : RegionName)
(chunk BLOCK_SIZE : Nat) :
varLenOneChunkIO old_a_start old_a_len old_a_location new_a_start
new_a_location chunk BLOCK_SIZE
⊨ fun _pid _len _olds _news xs i => xs i
IO definition: varLenOneChunkIOdef varLenOneChunkIO
(old_a_start old_a_len old_a_location new_a_start new_a_location : RegionName)
(chunk BLOCK_SIZE : Nat) : Meta3MaskedTileKernelIO₁ where
kernel := var_len_copy_one_chunk old_a_start old_a_len old_a_location
new_a_start new_a_location chunk BLOCK_SIZE
mbuf1 := old_a_len
mbuf2 := old_a_start
mbuf3 := new_a_start
inp := old_a_location
out := new_a_location
shape := [BLOCK_SIZE]
mwin1 := fun pid => pid
mwin2 := fun pid => pid
mwin3 := fun pid => pid
read := fun _pid _len olds _news i => olds + chunk * BLOCK_SIZE + i.1.val
write := fun _pid _len _olds news i => news + chunk * BLOCK_SIZE + i.1.val
mask := fun _pid len _olds _news i => i.1.val < len
Referenced model/specification definitions: Python source |
|---|
vector_additionadd_kernel_correctnessMathematical execution | Original-kernel candidate Masked per-program vector addition for arbitrary element count/block size with the declared IO windows; no host launch or concrete floating claim. | Statement and preconditionsspecification add_kernel_correctness
(x_ptr y_ptr output_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) :
addIO x_ptr y_ptr output_ptr n_elements BLOCK_SIZE
⊨ fun xs ys i => xs i + ys i
IO definition: addIOdef addIO (x_ptr y_ptr output_ptr : RegionName)
(n_elements BLOCK_SIZE : Nat) : MaskedKernelIO₂ where
kernel := add_kernel x_ptr y_ptr output_ptr n_elements BLOCK_SIZE
in1 := x_ptr
in2 := y_ptr
out := output_ptr
B := BLOCK_SIZE
read1 := fun pid => pid * BLOCK_SIZE
read2 := fun pid => pid * BLOCK_SIZE
write := fun pid => pid * BLOCK_SIZE
mask := fun pid j => pid * BLOCK_SIZE + j.val < n_elements
Referenced model/specification definitions: Python source |
|---|
vector_addition_customadd_kernel_correctnessMathematical execution | Original-kernel candidate Masked per-program custom vector addition for arbitrary size/block parameters with the declared IO windows; no host launch or concrete floating claim. | Statement and preconditionsspecification add_kernel_correctness
(A B C : RegionName)
(size BLOCK : Nat) :
addCustomIO A B C size BLOCK
⊨ fun as bs i => as i + bs i
IO definition: addCustomIOdef addCustomIO (A B C : RegionName)
(size BLOCK : Nat) : MaskedKernelIO₂ where
kernel := _add_kernel A B C size BLOCK
in1 := A
in2 := B
out := C
B := BLOCK
read1 := fun pid => pid * BLOCK
read2 := fun pid => pid * BLOCK
write := fun pid => pid * BLOCK
mask := fun pid j => pid * BLOCK + j.val < size
Referenced model/specification definitions: Python source |
|---|