From 0ca8ad221056731ec7d56fe715b1d8a285d5eab8 Mon Sep 17 00:00:00 2001 From: John Behm Date: Sat, 12 Sep 2026 19:03:24 +0200 Subject: [PATCH 1/3] feat: prove retained output workspace candidates (PS6136) --- .gitattributes | 1 + benchmarks/README.md | 9 + checks/ps6136.go | 801 ++++ checks/ps6136_abi_regression_test.go | 44 + checks/ps6136_analysis_fixture_test.go | 74 + checks/ps6136_analyzer.go | 130 + checks/ps6136_analyzer_test.go | 112 + checks/ps6136_backend.go | 248 ++ checks/ps6136_buffer_origin.go | 50 + checks/ps6136_compiler_test.go | 623 +++ checks/ps6136_compound_review_test.go | 107 + checks/ps6136_constructor_writes.go | 136 + checks/ps6136_constructors.go | 63 + checks/ps6136_constructors_test.go | 94 + checks/ps6136_consumer_walk.go | 30 + checks/ps6136_consumer_walk_test.go | 76 + checks/ps6136_contract_context.go | 123 + checks/ps6136_coverage.go | 143 + checks/ps6136_coverage_test.go | 77 + checks/ps6136_error_barrier.go | 158 + checks/ps6136_error_barrier_test.go | 73 + checks/ps6136_extents.go | 103 + checks/ps6136_extents_test.go | 67 + checks/ps6136_external_owner_test.go | 46 + checks/ps6136_factory.go | 80 + checks/ps6136_factory_error.go | 94 + checks/ps6136_field_effects.go | 104 + checks/ps6136_field_effects_test.go | 54 + checks/ps6136_forwarding.go | 172 + checks/ps6136_forwarding_test.go | 61 + checks/ps6136_head_owner_test.go | 77 + checks/ps6136_head_width.go | 168 + checks/ps6136_head_width_test.go | 46 + checks/ps6136_host_prefix.go | 152 + checks/ps6136_host_prefix_test.go | 118 + checks/ps6136_leaf_extent.go | 92 + checks/ps6136_leaf_extent_test.go | 114 + checks/ps6136_leaf_recorder.go | 80 + checks/ps6136_native_bindings.go | 374 ++ checks/ps6136_observation_invocation_test.go | 138 + checks/ps6136_observation_reads.go | 99 + checks/ps6136_owner_adversaries_test.go | 174 + checks/ps6136_owner_root.go | 189 + checks/ps6136_owner_root_test.go | 58 + checks/ps6136_owner_test.go | 42 + checks/ps6136_owner_transfers.go | 110 + checks/ps6136_owner_transfers_test.go | 52 + checks/ps6136_profile_recorder.go | 257 ++ checks/ps6136_profile_recorder_test.go | 102 + checks/ps6136_recorder_binding.go | 124 + checks/ps6136_recorder_binding_test.go | 67 + checks/ps6136_release.go | 94 + checks/ps6136_release_test.go | 96 + checks/ps6136_residency.go | 45 + checks/ps6136_residency_test.go | 37 + checks/ps6136_retention.go | 184 + checks/ps6136_returned_callable.go | 120 + checks/ps6136_returned_callable_test.go | 45 + checks/ps6136_selection.go | 112 + checks/ps6136_snapshot.go | 66 + checks/ps6136_snapshot_test.go | 79 + checks/ps6136_source_functions.go | 43 + checks/ps6136_struct_zero.go | 108 + checks/ps6136_struct_zero_test.go | 41 + checks/ps6136_test.go | 229 + checks/ps6136_write_nodes.go | 89 + checks/ps6136_zero_flow.go | 87 + checks/ps6136_zero_flow_test.go | 60 + .../ps6136-owner/after/decoder.go.txt | 3949 +++++++++++++++++ checks/testdata/ps6136-owner/after/gpt.go.txt | 404 ++ .../ps6136-owner/before/decoder.go.txt | 3896 ++++++++++++++++ .../testdata/ps6136-owner/before/gpt.go.txt | 389 ++ checks/testdata/ps6136-owner/cuda.go.txt | 1504 +++++++ .../testdata/ps6136-owner/export-api/main.go | 74 + checks/testdata/ps6136-owner/llamagpu.go.txt | 450 ++ checks/testdata/ps6136-owner/model-api.go.txt | 74 + checks/testdata/ps6136-owner/nn-api.go.txt | 15 + .../ps6136-owner/sampling_fastpath.go.txt | 135 + checks/testdata/ps6136-owner/vulkan.go.txt | 184 + checks/testdata/src/ps6136/gpt.go | 389 ++ config/config.go | 3 + config/output_workspace.go | 256 ++ config/output_workspace_test.go | 81 + docs/checks/PS6136.md | 44 + docs/checks/README.md | 3 +- docs/issue-887-output-workspace.md | 174 + internal/runner/ps6136_vocabulary_test.go | 16 + internal/runner/runner.go | 1 + 88 files changed, 20161 insertions(+), 1 deletion(-) create mode 100644 checks/ps6136.go create mode 100644 checks/ps6136_abi_regression_test.go create mode 100644 checks/ps6136_analysis_fixture_test.go create mode 100644 checks/ps6136_analyzer.go create mode 100644 checks/ps6136_analyzer_test.go create mode 100644 checks/ps6136_backend.go create mode 100644 checks/ps6136_buffer_origin.go create mode 100644 checks/ps6136_compiler_test.go create mode 100644 checks/ps6136_compound_review_test.go create mode 100644 checks/ps6136_constructor_writes.go create mode 100644 checks/ps6136_constructors.go create mode 100644 checks/ps6136_constructors_test.go create mode 100644 checks/ps6136_consumer_walk.go create mode 100644 checks/ps6136_consumer_walk_test.go create mode 100644 checks/ps6136_contract_context.go create mode 100644 checks/ps6136_coverage.go create mode 100644 checks/ps6136_coverage_test.go create mode 100644 checks/ps6136_error_barrier.go create mode 100644 checks/ps6136_error_barrier_test.go create mode 100644 checks/ps6136_extents.go create mode 100644 checks/ps6136_extents_test.go create mode 100644 checks/ps6136_external_owner_test.go create mode 100644 checks/ps6136_factory.go create mode 100644 checks/ps6136_factory_error.go create mode 100644 checks/ps6136_field_effects.go create mode 100644 checks/ps6136_field_effects_test.go create mode 100644 checks/ps6136_forwarding.go create mode 100644 checks/ps6136_forwarding_test.go create mode 100644 checks/ps6136_head_owner_test.go create mode 100644 checks/ps6136_head_width.go create mode 100644 checks/ps6136_head_width_test.go create mode 100644 checks/ps6136_host_prefix.go create mode 100644 checks/ps6136_host_prefix_test.go create mode 100644 checks/ps6136_leaf_extent.go create mode 100644 checks/ps6136_leaf_extent_test.go create mode 100644 checks/ps6136_leaf_recorder.go create mode 100644 checks/ps6136_native_bindings.go create mode 100644 checks/ps6136_observation_invocation_test.go create mode 100644 checks/ps6136_observation_reads.go create mode 100644 checks/ps6136_owner_adversaries_test.go create mode 100644 checks/ps6136_owner_root.go create mode 100644 checks/ps6136_owner_root_test.go create mode 100644 checks/ps6136_owner_test.go create mode 100644 checks/ps6136_owner_transfers.go create mode 100644 checks/ps6136_owner_transfers_test.go create mode 100644 checks/ps6136_profile_recorder.go create mode 100644 checks/ps6136_profile_recorder_test.go create mode 100644 checks/ps6136_recorder_binding.go create mode 100644 checks/ps6136_recorder_binding_test.go create mode 100644 checks/ps6136_release.go create mode 100644 checks/ps6136_release_test.go create mode 100644 checks/ps6136_residency.go create mode 100644 checks/ps6136_residency_test.go create mode 100644 checks/ps6136_retention.go create mode 100644 checks/ps6136_returned_callable.go create mode 100644 checks/ps6136_returned_callable_test.go create mode 100644 checks/ps6136_selection.go create mode 100644 checks/ps6136_snapshot.go create mode 100644 checks/ps6136_snapshot_test.go create mode 100644 checks/ps6136_source_functions.go create mode 100644 checks/ps6136_struct_zero.go create mode 100644 checks/ps6136_struct_zero_test.go create mode 100644 checks/ps6136_test.go create mode 100644 checks/ps6136_write_nodes.go create mode 100644 checks/ps6136_zero_flow.go create mode 100644 checks/ps6136_zero_flow_test.go create mode 100644 checks/testdata/ps6136-owner/after/decoder.go.txt create mode 100644 checks/testdata/ps6136-owner/after/gpt.go.txt create mode 100644 checks/testdata/ps6136-owner/before/decoder.go.txt create mode 100644 checks/testdata/ps6136-owner/before/gpt.go.txt create mode 100644 checks/testdata/ps6136-owner/cuda.go.txt create mode 100644 checks/testdata/ps6136-owner/export-api/main.go create mode 100644 checks/testdata/ps6136-owner/llamagpu.go.txt create mode 100644 checks/testdata/ps6136-owner/model-api.go.txt create mode 100644 checks/testdata/ps6136-owner/nn-api.go.txt create mode 100644 checks/testdata/ps6136-owner/sampling_fastpath.go.txt create mode 100644 checks/testdata/ps6136-owner/vulkan.go.txt create mode 100644 checks/testdata/src/ps6136/gpt.go create mode 100644 config/output_workspace.go create mode 100644 config/output_workspace_test.go create mode 100644 docs/checks/PS6136.md create mode 100644 docs/issue-887-output-workspace.md create mode 100644 internal/runner/ps6136_vocabulary_test.go diff --git a/.gitattributes b/.gitattributes index 3cfc61e8..0ff85f86 100644 --- a/.gitattributes +++ b/.gitattributes @@ -34,3 +34,4 @@ checks/testdata/ps6133_owner_vulkan_shader.txt text eol=lf checks/testdata/src/ps6133/metal_layout.h text eol=lf checks/testdata/src/ps6133/vulkan_layout.h text eol=lf checks/testdata/ps6137_owner_*.txt text eol=lf +checks/testdata/ps6136-owner/**/*.txt text eol=lf diff --git a/benchmarks/README.md b/benchmarks/README.md index d7533e44..60217f82 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -279,6 +279,15 @@ Rules without a pair, and why: measurements involve complete scratch/backend changes, not isolated binElem fallback support. No PS6135-specific speedup is claimed; preserve every shape/error/overflow/tail/provider path and measure the complete operation. +- **PS6136** (maximum-context retained output workspace advisory): symbolic + residency is source-bound constructor geometry under a reviewed one-row + versus rare bulk policy, not a native throughput prediction. Frozen actual + GPT/Llama before/after owner replay and adversarial source/lifetime tests + cannot execute kernels or demonstrate allocator/completion ownership on + every provider. Owner PR1207 measurements are attributed complete-public-API + campaigns, not a portable toy benchmark or an isolated detector speedup. + Preserve physical ToHost transfer, all native parity/errors, checked shapes + and release/publication; remeasure residency and complete public workloads. PS6137 is an opt-in advisory with no automatic replacement benchmark. Its complete owner cgo/source replay is not a native workload execution. GoAI diff --git a/checks/ps6136.go b/checks/ps6136.go new file mode 100644 index 00000000..fd2e3c62 --- /dev/null +++ b/checks/ps6136.go @@ -0,0 +1,801 @@ +package checks + +import ( + "go/ast" + "go/token" + "go/types" + "maps" + + "github.com/jxsl13/perfscan/config" + "golang.org/x/tools/go/analysis" + "golang.org/x/tools/go/ssa" +) + +// The assembled selection proof keeps source identities separate from the +// contract's reviewed model/backend/workload facts. A selection alone is not +// reportable: complete observations and constructor/lifetime gates follow it. +type ps6136Selection struct { + context *ps6125SSAContext + owner ps6125SSAReference + model ps6125SSAReference + ownerType, modelType *types.Named + maximumRows, width, workspace, projector, ops, retained, slot *types.Var + modelConfig, configRows, configWidth, modelHead *types.Var + allocator *types.Var + primaryRecorder, secondaryRecorder *types.Var + allocatorIndex int + allocatorIDs map[string]bool + allocationFunction, releaseMethod, shapeMethod string + retainedReleaseMethod string + projectorWidths map[*types.Named]*types.Var +} + +func (selection *ps6136Selection) closedSource(pass *analysis.Pass, pkg *ssa.Package, consumers []*ps6136ConsumerProof) bool { + if selection == nil || pass == nil || pkg == nil { + return false + } + fields := map[*types.Var]bool{selection.workspace: true, selection.width: true, selection.maximumRows: true, selection.ops: true, selection.projector: true, selection.retained: true} + writes := ps6136FreshConstructorWrites(pkg, selection.ownerType, fields) + if writes == nil { + return false + } + for _, field := range writes.readFields { + if field == selection.workspace { + return false + } + } + initializingWrites := *writes + initializingWrites.writers = maps.Clone(writes.writers) + for writer := range initializingWrites.writers { + if object, ok := writer.Object().(*types.Func); ok && ps6090FunctionID(object) == selection.releaseMethod { + delete(initializingWrites.writers, writer) + } + } + initializers := ps6136ConstructorWriteNodes(pass, pkg, selection.ownerType, fields, &initializingWrites) + if initializers == nil { + return false + } + list, ok := selection.retained.Type().Underlying().(*types.Slice) + if !ok { + return false + } + bufferRelease, _, _ := types.LookupFieldOrMethod(list.Elem(), true, pass.Pkg, "Release") + release, releaseOK := bufferRelease.(*types.Func) + if !releaseOK { + return false + } + if selection.retainedReleaseMethod == "" || ps6090FunctionID(release) != selection.retainedReleaseMethod { + return false + } + var listRelease map[ast.Node]bool + for _, file := range pass.Files { + for _, declaration := range file.Decls { + function, ok := declaration.(*ast.FuncDecl) + if !ok || function.Recv == nil || len(function.Recv.List) != 1 || len(function.Recv.List[0].Names) != 1 { + continue + } + object, ok := pass.TypesInfo.Defs[function.Name].(*types.Func) + if !ok || ps6090FunctionID(object) != selection.releaseMethod { + continue + } + listRelease = ps6136ReleaseList(pass, function, pass.TypesInfo.Defs[function.Recv.List[0].Names[0]], selection.retained, release) + } + } + if listRelease == nil { + return false + } + for node := range listRelease { + initializers[node] = true + } + for field := range fields { + if !ps6136FieldEffects(pass, field, initializers) { + return false + } + } + // Value-literal keys initialize new backendOps values, not a published + // owner's callback cell. All subsequent direct callback writes/addresses + // remain forbidden, independently of sibling recorder policy overrides. + allocatorInitializers := make(map[ast.Node]bool) + for _, file := range pass.Files { + ast.Inspect(file, func(node ast.Node) bool { + if pair, ok := node.(*ast.KeyValueExpr); ok { + if key, ok := pair.Key.(*ast.Ident); ok && pass.TypesInfo.Uses[key] == selection.allocator { + allocatorInitializers[pair] = true + } + } + return true + }) + } + if !ps6136FieldEffects(pass, selection.allocator, allocatorInitializers) { + return false + } + if !ps6136FieldEffects(pass, selection.modelConfig, nil) || !ps6136FieldEffects(pass, selection.modelHead, nil) { + return false + } + proved := make(map[ast.Node]bool) + for node := range initializers { + if pair, ok := node.(*ast.KeyValueExpr); ok { + proved[pair.Key] = true + } else { + proved[node] = true + } + } + for _, file := range pass.Files { + ast.Inspect(file, func(node ast.Node) bool { + assignment, ok := node.(*ast.AssignStmt) + if !ok || len(assignment.Lhs) != 1 || len(assignment.Rhs) != 1 || !initializers[assignment.Lhs[0]] { + return true + } + target, ok := assignment.Lhs[0].(*ast.SelectorExpr) + if !ok { + return true + } + selectionInfo := pass.TypesInfo.Selections[target] + owner, ok := target.X.(*ast.Ident) + if !ok || selectionInfo == nil || selectionInfo.Obj() != selection.retained { + return true + } + call, ok := assignment.Rhs[0].(*ast.CallExpr) + if !ok || len(call.Args) != 2 || call.Ellipsis.IsValid() { + return true + } + builtin, ok := call.Fun.(*ast.Ident) + if !ok || pass.TypesInfo.Uses[builtin] != types.Universe.Lookup("append") || !ps6136DirectField(pass, call.Args[0], pass.TypesInfo.Uses[owner], selection.retained) { + return true + } + proved[call.Args[0]] = true + return true + }) + } + // readFields and positions use sparse offsets across independent owner files. + positions := make(map[token.Pos]bool) + calls := make(map[*ssa.Call]bool, len(writes.calls)) + for position := range writes.reads { + if writes.readFields[position] != selection.retained { //perfscan:ignore PS3003 + positions[position] = true //perfscan:ignore PS3003 + } + } + for call := range writes.calls { + calls[call] = true + } + for _, consumer := range consumers { + if consumer == nil { + return false + } + for position := range consumer.reads { + positions[position] = true + } + for call := range consumer.calls { + calls[call] = true + } + } + constructors := ps6136ConstructorInventory(pkg, selection.ownerType) + if constructors == nil { + return false + } + for _, file := range pass.Files { + listClosed := true + ast.Inspect(file, func(node ast.Node) bool { + switch occurrence := node.(type) { + case *ast.SelectorExpr: + selectionInfo := pass.TypesInfo.Selections[occurrence] + if selectionInfo != nil && selectionInfo.Obj() == selection.retained { + if !proved[occurrence] { + for position := range positions { + if occurrence.Pos() <= position && position < occurrence.End() { + proved[occurrence] = true + } + } + } + if !proved[occurrence] { + listClosed = false + } + } + if selectionInfo == nil || selectionInfo.Obj() != selection.workspace { + break + } + for position := range positions { + if occurrence.Pos() <= position && position < occurrence.End() { + proved[occurrence] = true + } + } + case *ast.CallExpr: + function, _, typed := typedCallee(pass, occurrence.Fun) + if !typed { + break + } + approved := false + for call := range calls { + callee := call.Call.StaticCallee() + if callee != nil && callee.Object() == function && occurrence.Pos() <= call.Pos() && call.Pos() < occurrence.End() { + approved = true + } + } + if !approved { + break + } + if selector, ok := ps2110Unparen(occurrence.Fun).(*ast.SelectorExpr); ok { + if identifier, ok := ps2110Unparen(selector.X).(*ast.Ident); ok && types.Identical(pass.TypesInfo.TypeOf(identifier), types.NewPointer(selection.ownerType)) { + proved[identifier] = true + } + } + for _, argument := range occurrence.Args { + if identifier, ok := ps2110Unparen(argument).(*ast.Ident); ok && types.Identical(pass.TypesInfo.TypeOf(identifier), types.NewPointer(selection.ownerType)) { + proved[identifier] = true + } + } + case *ast.FuncDecl: + object, ok := pass.TypesInfo.Defs[occurrence.Name].(*types.Func) + if !ok { + break + } + function := pkg.Prog.FuncValue(object) + if constructors[function].value == nil { + break + } + ast.Inspect(occurrence.Body, func(node ast.Node) bool { + if _, nested := node.(*ast.FuncLit); nested { + return false + } + if returned, ok := node.(*ast.ReturnStmt); ok && len(returned.Results) != 0 { + if identifier, ok := ps2110Unparen(returned.Results[0]).(*ast.Ident); ok && types.Identical(pass.TypesInfo.TypeOf(identifier), types.NewPointer(selection.ownerType)) { + proved[identifier] = true + } + } + return true + }) + } + return true + }) + if !listClosed { + return false + } + } + return ps6136ClosedObservations(pass, selection.ownerType, selection.workspace, proved) +} + +type ps6136ConsumerProof struct { + reads map[token.Pos]bool + calls map[*ssa.Call]bool + capacityTransfers int + boundedLeaves int +} + +func (selection *ps6136Selection) consumerCalls(context *ps6125SSAContext, leaves []config.OutputWorkspaceLeaf, bulk bool, capacity func(*ps6125SSAContext, *ssa.Call, *config.OutputWorkspaceLeaf) bool) *ps6136ConsumerProof { + if selection == nil || context == nil || len(context.flow.function.Params) == 0 { + return nil + } + parameter := context.flow.function.Params[0] + if !types.Identical(parameter.Type(), types.NewPointer(selection.ownerType)) || parameter.Object() == nil { + return nil + } + owner := context.reference(parameter) + var pool ps6125ExtentPool + width := ps6125SymbolicExtent(pool.identity(parameter.Object(), []*types.Var{selection.width}, false)) + var bulkExtent ps6125Extent + if bulk { + if len(context.flow.function.Params) < 2 || !types.Identical(context.flow.function.Params[1].Type(), types.NewSlice(types.Typ[types.Int])) { + return nil + } + tokens := context.flow.function.Params[1] + length := ps6125SymbolicExtent(pool.identity(tokens.Object(), nil, true)) + // Recompute the contextual scalar/branch algebra after binding the + // actual bulk descriptor, retaining the same invocation identity. + context.flow = ps6125AnalyzeSSAExtents(context.flow.function, nil, map[*ssa.Parameter]ps6125Extent{tokens: length}) + bulkExtent = ps6125MultiplyExtents(width, length) + } + facts := &ps6136Extents{owner: owner, ownerType: selection.ownerType, immutable: map[*types.Var]ps6125Extent{selection.width: width}, remaining: 16384} + proof := &ps6136ConsumerProof{reads: make(map[token.Pos]bool), calls: make(map[*ssa.Call]bool)} + leafCalls := make(map[*ps6125SSAContext]map[*ssa.Call]bool) + readValues := make(map[ps6125SSAReference]token.Pos) + valid := true + bulkWitness := false + remaining := 16384 + if !ps6136WalkConsumerCalls(context, func(current *ps6125SSAContext, call *ssa.Call) bool { + if !valid { + return true + } + paths := ps6125AccessPaths{flow: current.flow} + for _, block := range current.flow.function.Blocks { + if !current.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + load, ok := instruction.(*ssa.UnOp) + if !ok || load.Op != token.MUL { + continue + } + path := paths.resolve(load.X) + if !path.known || len(path.access.fields) != 1 || path.access.fields[0] != selection.workspace { + continue + } + if ps6136AccessOwnerRoot(current, path, selection.ownerType) != owner { + valid = false + return true + } + readValues[ps6125SSAReference{context: current, value: load}] = load.Pos() + } + } + for index := range leaves { + leaf := &leaves[index] + extent, ok := ps6136LeafExtent(current, call, leaf, facts, selection.workspace, selection.slot, selection.projector, selection.width) + if !ok { + continue + } + if (leaf.Kind == "projection" || leaf.Kind == "row-width") && !selection.leafRecorderOrigin(current, call, leaf.Kind, owner) { + valid = false + return true + } + if leaf.Kind == "capacity-transfer" { + if capacity == nil || !capacity(current, call, leaf) { + valid = false + return true + } + proof.capacityTransfers++ + } else { + if !ps6125SameExtent(extent, width) && (!bulk || !ps6125SameExtent(extent, bulkExtent)) { + valid = false + return true + } + proof.boundedLeaves++ + if bulk && ps6125SameExtent(extent, bulkExtent) { + bulkWitness = true + } + } + proof.calls[call] = true + if leafCalls[current] == nil { + leafCalls[current] = make(map[*ssa.Call]bool) + } + leafCalls[current][call] = true + return true + } + child := ps6136Call(current, call) + for _, argument := range call.Call.Args { + if types.Identical(argument.Type(), types.NewPointer(selection.ownerType)) { + if child == nil || len(child.flow.function.Blocks) == 0 || ps6136OwnerRoot(current, argument, selection.ownerType) != owner { + valid = false + return true + } + proof.calls[call] = true + } + } + return false + }, &remaining) || !valid || context.budget.remaining == 0 { + return nil + } + for reference, position := range readValues { + // Only separately proved leaves, not ordinary helper edges, may + // terminate the workspace's buffer-alias use graph. + if !position.IsValid() || !ps6136WorkspaceReadClosed(reference.value, selection.slot, leafCalls[reference.context], 512) { + return nil + } + proof.reads[position] = true + } + if proof.boundedLeaves == 0 || bulk && !bulkWitness { + return nil + } + return proof +} + +func (selection *ps6136Selection) configValue(context *ps6125SSAContext, value ssa.Value, wanted *types.Var, budget int) bool { + if selection == nil || context == nil || value == nil || wanted == nil || budget <= 0 || !types.Identical(value.Type(), types.Typ[types.Int]) { + return false + } + path := ps6136SelectionPath(context, value, budget) + return path.known && len(path.fields) == 2 && path.fields[0] == selection.modelConfig && path.fields[1] == wanted && selection.model == path.root +} + +type ps6136SelectionAccess struct { + root ps6125SSAReference + fields []*types.Var + known bool +} + +// Compose typed field paths through actual invocation bindings. Loading a +// config value is only provenance, not a claim its mutable model stayed stable. +func ps6136SelectionPath(context *ps6125SSAContext, value ssa.Value, budget int) ps6136SelectionAccess { + if context == nil || value == nil || budget <= 0 { + return ps6136SelectionAccess{} + } + reference := context.reference(value) + if reference.value == nil { + return ps6136SelectionAccess{} + } + if reference.context != context || reference.value != value { + return ps6136SelectionPath(reference.context, reference.value, budget-1) + } + paths := ps6125AccessPaths{flow: context.flow} + path := paths.resolve(value) + if !path.known { + return ps6136SelectionAccess{} + } + root := context.reference(path.access.root) + if root.value == nil { + return ps6136SelectionAccess{} + } + if allocation, ok := root.value.(*ssa.Alloc); ok { + stores := root.context.structCell(allocation) + point, ok := value.(ssa.Instruction) + if stores == nil || stores.whole == nil || !ok || root.context != context || !context.flow.instructionDominates(stores.whole, point) { + return ps6136SelectionAccess{} + } + prefix := ps6136SelectionPath(context, stores.whole.Val, budget-1) + if !prefix.known { + return prefix + } + prefix.fields = append(prefix.fields, path.access.fields...) + return prefix + } + if root.context != context || root.value != path.access.root { + prefix := ps6136SelectionPath(root.context, root.value, budget-1) + if !prefix.known { + return prefix + } + prefix.fields = append(prefix.fields, path.access.fields...) + return prefix + } + return ps6136SelectionAccess{root: root, fields: path.access.fields, known: true} +} + +func (selection *ps6136Selection) initializationValues(remaining int) bool { + if selection == nil || selection.context == nil || selection.owner.value == nil || selection.model.value == nil { + return false + } + counts := make(map[*types.Var]int) + pending := []*ps6125SSAContext{selection.context} + for len(pending) != 0 { + context := pending[len(pending)-1] + pending = pending[:len(pending)-1] + remaining-- + if remaining <= 0 { + return false + } + paths := ps6125AccessPaths{flow: context.flow} + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + if store, ok := instruction.(*ssa.Store); ok { + path := paths.resolve(store.Addr) + if path.known && len(path.access.fields) == 1 && ps6136AccessOwnerRoot(context, path, selection.ownerType) == selection.owner { + field := path.access.fields[0] + if field == selection.maximumRows || field == selection.width { + wanted := selection.configWidth + if field == selection.maximumRows { + wanted = selection.configRows + } + if !selection.configValue(context, store.Val, wanted, 128) { + return false + } + counts[field]++ + } + } + } + if call, ok := instruction.(*ssa.Call); ok { + if child := ps6136Call(context, call); child != nil && len(child.flow.function.Blocks) != 0 { + pending = append(pending, child) + } + } + if store, ok := instruction.(*ssa.UnOp); ok && store.Op == token.ARROW { + return false + } + } + } + } + return counts[selection.maximumRows] == 1 && counts[selection.width] == 1 +} + +type ps6136ConstructorProof struct { + allocation *ssa.Store + factory *ps6125SSAContext + errorCell ps6125SSAReference + residency ps6136Residency +} + +func (selection *ps6136Selection) constructorAllocation(remaining int) *ps6136ConstructorProof { + if selection == nil || !selection.initializationValues(remaining) || selection.workspace == nil || selection.allocator == nil || selection.ops == nil || selection.slot == nil || selection.retained == nil || selection.projector == nil { + return nil + } + model, ok := selection.model.value.(*ssa.Parameter) + if !ok || model.Object() == nil { + return nil + } + var pool ps6125ExtentPool + rows := pool.identity(model.Object(), []*types.Var{selection.modelConfig, selection.configRows}, false) + width := pool.identity(model.Object(), []*types.Var{selection.modelConfig, selection.configWidth}, false) + residency, ok := ps6136OutputResidency(rows, width) + if !ok { + return nil + } + facts := &ps6136Extents{owner: selection.owner, ownerType: selection.ownerType, immutable: map[*types.Var]ps6125Extent{selection.maximumRows: ps6125SymbolicExtent(rows), selection.width: ps6125SymbolicExtent(width)}, remaining: remaining} + proof := &ps6136ConstructorProof{residency: residency} + pending := []*ps6125SSAContext{selection.context} + headStores := 0 + for len(pending) != 0 { + context := pending[len(pending)-1] + pending = pending[:len(pending)-1] + remaining-- + if remaining <= 0 { + return nil + } + paths := ps6125AccessPaths{flow: context.flow} + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + store, ok := instruction.(*ssa.Store) + if !ok { + continue + } + path := paths.resolve(store.Addr) + if !path.known || len(path.access.fields) != 1 || ps6136AccessOwnerRoot(context, path, selection.ownerType) != selection.owner { + continue + } + switch path.access.fields[0] { + case selection.projector: + if !ps6136ProjectorHeadWidth(context, store.Val, selection.projectorWidths, selection.model, selection.modelType, selection.modelHead, selection.shapeMethod, 1, 512) { + return nil + } + headStores++ + case selection.workspace: + function, ok := context.flow.function.Object().(*types.Func) + if proof.allocation != nil || !ok || ps6090FunctionID(function) != selection.allocationFunction { + return nil + } + call, ok := store.Val.(*ssa.Call) + if !ok || len(call.Call.Args) != 1 { + return nil + } + input := context.reference(call.Call.Args[0]) + allocation, ok := input.value.(*ssa.MakeSlice) + if !ok || !types.Identical(allocation.Type(), types.NewSlice(types.Typ[types.Float32])) || !ps6125SameExtent(facts.length(input.context, allocation), ps6125MultiplyExtents(ps6125SymbolicExtent(rows), ps6125SymbolicExtent(width))) { + return nil + } + factory := ps6136Call(context, call) + backend := ps6136FactoryResult(factory, input, selection.allocator, selection.slot) + if backend == nil || !ps6136FactoryRetention(factory, backend, selection.owner, selection.ownerType, selection.retained, selection.slot) { + return nil + } + if factory.call(backend) == nil { + snapshot := ps6136InitialOwnerField(selection.owner, selection.ops) + if snapshot.value == nil { + return nil + } + origin := snapshot.context.structField(snapshot.value, selection.allocatorIndex) + if ps6136CallOrigin(factory, backend, origin) == nil { + return nil + } + } + if !ps6136BoundAllocator(factory, backend, input, selection.allocatorIDs) { + return nil + } + errorCell := ps6136FactoryErrorCell(factory, backend) + if errorCell.value == nil || !ps6136ConstructorErrorBarrier(selection.context, selection.owner, selection.ownerType, errorCell, selection.releaseMethod) { + return nil + } + proof.allocation, proof.factory, proof.errorCell = store, factory, errorCell + } + } + } + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + if call, ok := instruction.(*ssa.Call); ok { + if child := ps6136Call(context, call); child != nil && len(child.flow.function.Blocks) != 0 { + pending = append(pending, child) + } + } + } + } + } + if proof.allocation == nil || headStores != 1 { + return nil + } + return proof +} + +func (selection *ps6136Selection) constructorLifetime(proof *ps6136ConstructorProof, remaining int) bool { + if selection == nil || proof == nil || proof.errorCell.value == nil { + return false + } + pending := []*ps6125SSAContext{selection.context} + phase := make(map[ssa.Instruction]bool) + for len(pending) != 0 { + context := pending[len(pending)-1] + pending = pending[:len(pending)-1] + remaining-- + if remaining <= 0 { + return false + } + paths := ps6125AccessPaths{flow: context.flow} + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + switch use := instruction.(type) { + case *ssa.Defer, *ssa.Go: + return false // deferred allocation/error effects can run after publication + case *ssa.Store: + if ps6136ErrorCellRoot(context, use.Addr) == proof.errorCell { + errorValue, ok := use.Val.(*ssa.Extract) + if !ok || errorValue.Index != 1 { + return false + } + backend, ok := errorValue.Tuple.(*ssa.Call) + if !ok || ps6136FactoryErrorCell(context, backend) != proof.errorCell { + return false + } + } + if ps6136ErrorCellRoot(context, use.Val) == proof.errorCell { + allocation, ok := use.Addr.(*ssa.Alloc) + if !ok || ps6136ClosedOwnerCell(context, allocation) != use { + return false + } + } + if _, function := use.Val.Type().Underlying().(*types.Signature); function { + allocation, ok := use.Addr.(*ssa.Alloc) + if !ok || context.functionCell(allocation) != use { + return false + } + } + case *ssa.Convert: + if ps6136ErrorCellRoot(context, use.X) == proof.errorCell { + return false + } + case *ssa.MakeInterface: + if ps6136ErrorCellRoot(context, use.X) == proof.errorCell { + return false + } + case *ssa.Call: + path := paths.resolve(use.Call.Value) + if path.known && len(path.access.fields) != 0 && path.access.fields[len(path.access.fields)-1] == selection.allocator { + if len(use.Call.Args) != 1 { + return false + } + input := context.reference(use.Call.Args[0]) + if ps6136FactoryResult(context, input, selection.allocator, selection.slot) != use || ps6136FactoryErrorCell(context, use) != proof.errorCell || !ps6136FactoryRetention(context, use, selection.owner, selection.ownerType, selection.retained, selection.slot) { + return false + } + if context.call(use) == nil { + snapshot := ps6136InitialOwnerField(selection.owner, selection.ops) + if snapshot.value == nil { + return false + } + if ps6136CallOrigin(context, use, snapshot.context.structField(snapshot.value, selection.allocatorIndex)) == nil { + return false + } + } + if !ps6136BoundAllocator(context, use, input, selection.allocatorIDs) { + return false + } + var point ssa.Instruction = use + ancestor := context + for ancestor != selection.context && ancestor.parent != nil { + point = ancestor.site + ancestor = ancestor.parent + } + if ancestor != selection.context || point == nil { + return false + } + phase[point] = true + } + child := ps6136Call(context, use) + for _, argument := range use.Call.Args { + _, function := argument.Type().Underlying().(*types.Signature) + if function || ps6136ErrorCellRoot(context, argument) == proof.errorCell || types.Identical(argument.Type(), types.NewPointer(selection.ownerType)) { + if child == nil || len(child.flow.function.Blocks) == 0 { + return false + } + } + } + if child != nil && len(child.flow.function.Blocks) != 0 { + ownerResult := false + for index := 0; index < child.flow.function.Signature.Results().Len(); index++ { + if types.Identical(child.flow.function.Signature.Results().At(index).Type(), types.NewPointer(selection.ownerType)) { + ownerResult = true + } + } + if !ownerResult { + child = ps6136ZeroSourceFlow(child) + } + pending = append(pending, child) + } + } + } + } + } + if len(phase) == 0 { + return false + } + for _, block := range selection.context.flow.function.Blocks { + if !selection.context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + returned, ok := instruction.(*ssa.Return) + if !ok || len(returned.Results) < 1 { + continue + } + constant, nilOwner := returned.Results[0].(*ssa.Const) + if !nilOwner || !constant.IsNil() { + continue + } + afterAllocation := false + for point := range phase { + if ps6136InstructionCanPrecede(selection.context, point, returned) { + afterAllocation = true + } + } + if !afterAllocation { + continue + } + cleanup := false + for _, call := range ps6136Calls(selection.context) { + callee := call.Call.StaticCallee() + if callee == nil || callee.Object() == nil || ps6090FunctionID(callee.Object().(*types.Func)) != selection.releaseMethod || len(call.Call.Args) != 1 { + continue + } + if ps6136OwnerRoot(selection.context, call.Call.Args[0], selection.ownerType) == selection.owner && selection.context.flow.instructionDominates(call, returned) { + cleanup = true + } + } + if !cleanup { + return false + } + } + } + return true +} + +func ps6136Calls(context *ps6125SSAContext) []*ssa.Call { + var result []*ssa.Call + for _, block := range context.flow.function.Blocks { + if context.flow.blocks[block] { + for _, instruction := range block.Instrs { + if call, ok := instruction.(*ssa.Call); ok { + result = append(result, call) + } + } + } + } + return result +} + +func ps6136InstructionCanPrecede(context *ps6125SSAContext, first, last ssa.Instruction) bool { + if first.Block() == last.Block() { + for _, instruction := range first.Block().Instrs { + if instruction == first { + return true + } + if instruction == last { + return false + } + } + return false + } + pending := []*ssa.BasicBlock{first.Block()} + seen := make(map[*ssa.BasicBlock]bool) + for len(pending) != 0 { + block := pending[len(pending)-1] + pending = pending[:len(pending)-1] + if seen[block] { + continue + } + seen[block] = true + for _, next := range block.Succs { + if !context.flow.edges[ps6125SSAEdge{from: block, to: next}] { + continue + } + if next == last.Block() { + return true + } + pending = append(pending, next) + } + } + return false +} diff --git a/checks/ps6136_abi_regression_test.go b/checks/ps6136_abi_regression_test.go new file mode 100644 index 00000000..af1251ee --- /dev/null +++ b/checks/ps6136_abi_regression_test.go @@ -0,0 +1,44 @@ +package checks + +import ( + "go/token" + "go/types" + "golang.org/x/tools/go/analysis" + "testing" +) + +func TestPS6136DottedPackageFunctionCannotReplaceMethod(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, "before", true, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg} + context := ps6136ContractsContext(pass) + c := ps6136CompleteOwnerContract(pkg, "GPTDecoder") + selection := context.selection(pkg, context.callable(c.ConstructorEntries[0]), &c) + if selection == nil { + t.Fatal("actual owner selection missing") + } + proof := selection.constructorAllocation(16384) + if proof == nil { + t.Fatal("actual allocation missing") + } + dotted := types.NewPackage("example.test/native.v1", "v1") + signature := types.NewSignatureType(nil, nil, nil, types.NewTuple(), types.NewTuple(types.NewVar(token.NoPos, dotted, "err", types.Universe.Lookup("error").Type())), false) + function := types.NewFunc(token.NoPos, dotted, "Project", signature) + dotted.Scope().Insert(function) + dotted.MarkComplete() + context.packages[dotted.Path()] = dotted + c.ProjectorRecordMethod = dotted.Path() + ".Project" + if !c.Valid() { + t.Fatal("receiver-looking package suffix does not reproduce valid syntax") + } + if context.callable(c.ProjectorRecordMethod) != function { + t.Fatal("dotted package function did not resolve") + } + if context.nativeBindings(selection, proof, &c) { + t.Fatal("package function accepted as required projector method") + } +} diff --git a/checks/ps6136_analysis_fixture_test.go b/checks/ps6136_analysis_fixture_test.go new file mode 100644 index 00000000..509088ab --- /dev/null +++ b/checks/ps6136_analysis_fixture_test.go @@ -0,0 +1,74 @@ +package checks + +import ( + "bytes" + "fmt" + "go/ast" + "go/format" + "go/parser" + "go/token" + "os" + "path/filepath" + "testing" + + "github.com/jxsl13/perfscan/config" + "golang.org/x/tools/go/analysis" + "golang.org/x/tools/go/analysis/analysistest" +) + +// Conventional analysistest coverage uses the entire pinned GPT owner fixture, +// together with the same complete constructor/observation census as the replay. +// Temporary dependency files are explicitly type-only native/model scaffolds. +func TestPS6136AnalysisFixture(t *testing.T) { + t.Parallel() + data, err := os.ReadFile("testdata/src/ps6136/gpt.go") + if err != nil { + t.Fatal(err) + } + original, err := os.ReadFile("testdata/ps6136-owner/before/gpt.go.txt") + if err != nil { + t.Fatal(err) + } + withoutExpectation := bytes.Replace(data, []byte(` // want "reviewed common one-row output policy retains"`), nil, 1) + if !bytes.Equal(withoutExpectation, original) { + t.Fatal("conventional fixture differs from complete pinned owner source") + } + fixture, err := ps6136CompileOwnersRewrite(t, "before", true, true, func(fs *token.FileSet, files []*ast.File) []*ast.File { + file, err := parser.ParseFile(fs, "testdata/src/ps6136/gpt.go", nil, parser.ParseComments|parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + files[0] = file + return files + }) + if err != nil { + t.Fatal(err) + } + c := ps6136CompleteOwnerContract(ps6136FixtureSSA(fixture), "GPTDecoder") + dir := t.TempDir() + write := func(path string, data []byte) { + t.Helper() + path = filepath.Join(dir, "src", path) + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, data, 0o600); err != nil { + t.Fatal(err) + } + } + for path, source := range fixture.sources { + write(filepath.Join(path, "api.go"), []byte(source)) + } + for index, file := range fixture.files { + var source bytes.Buffer + if err := format.Node(&source, fixture.fileset, file); err != nil { + t.Fatal(err) + } + write(filepath.Join(fixture.pkg.Path(), fmt.Sprintf("owner%d.go", index)), source.Bytes()) + } + analyzer := *PS6136.Analyzer + analyzer.Run = func(pass *analysis.Pass) (any, error) { + return runPS6136WithContracts(pass, []config.OutputWorkspaceContract{c}) + } + analysistest.Run(t, dir, &analyzer, fixture.pkg.Path()) +} diff --git a/checks/ps6136_analyzer.go b/checks/ps6136_analyzer.go new file mode 100644 index 00000000..2f6b70ed --- /dev/null +++ b/checks/ps6136_analyzer.go @@ -0,0 +1,130 @@ +package checks + +import ( + "github.com/jxsl13/perfscan/config" + "github.com/jxsl13/perfscan/lint" + "go/ast" + "go/token" + "go/types" + "golang.org/x/tools/go/analysis" + "golang.org/x/tools/go/ssa" + "math" + "slices" + "strconv" +) + +var PS6136 = register(&lint.Check{ + ID: "PS6136", Category: "alloc", Slug: "maximum-context-retained-output-workspace", Level: lint.LevelAggressive, + AutoFix: false, NeedsConfig: true, Vocab: []string{"outputWorkspaceContracts"}, + Doc: lint.Documentation{ + Title: "maximum-context output scratch remains resident for a reviewed one-row inference policy", + Text: "PS6136 requires an exact typed outputWorkspaceContracts entry and proves a selected public model constructor's source flow into a fresh owner, context-times-width float32 output allocation, exact owning backend factory/retention list, current final error barrier, source head Shape width, complete loaded-package constructor/field-use inventory, and selected-instance producer/consumer extents. It does not identify models, projectors or allocators from names.\n\nThe selected common methods must consume one logical output row; an explicit rare bulk method may consume actual token rows. Whole-capacity ToHost is inventoried separately as a physical transfer, with a source-proved first-width logical prefix; it is never counted as a one-row physical leaf. Exact typed selected native buffer/recorder factories, wrappers, forwarding and capability availability are independently bound. Unknown/unstable arguments, aliases, unclassified observations, allocator/geometry changes, unsafe constructor error publication, full logical consumption, contradictory contracts and hard-real-time no-growth policy remain silent.\n\nThe contract supplies the narrowly reviewed model Config width/head relationship, native allocation/access/release/completion semantics, positive checked shapes, provider/build partitions and external observers, sequential ownership and dominant one-row versus rare bulk policy. These are not proved by Go syntax, a source hash or a method name. Unselected architecture/model factories remain outside the selected model-width contract.\n\nThere is no automatic fix. Review one resident output row plus a privately built exact high-water bulk allocation. Preserve all public hot/bulk paths, physical transfer and synchronization, invalid shapes, numeric/reference and sequential/batched parity, partial failure cleanup, final release and publication ownership. A hard-real-time allocator policy may deliberately retain maximum capacity. Measure retained bytes and public constructor/Step/StepNLast/StepN campaigns in order-alternating same-binary runs; do not attribute unrelated projector/backend work to storage alone.", + Before: "d.logits = mk(make([]float32, d.maxLen*d.v))\n// Common logical output: d.v; explicit rare bulk: rows*d.v.", + After: "// Keep one resident row; privately allocate/reuse a checked exact bulk generation;\n// publish only after success; preserve all reads, failures, completion and release.", + MeasuredWin: "Owner issue #887 / GoAI PR #1207 report GPT-2 small context 1024, vocabulary 50257: 205852672 retained float32 output bytes versus 201028 one-row bytes, 205651644 idle bytes and 1024x amplification. Apple M2 Pro same-binary constructor medians were 3423267 versus 157542 ns (21.73x), with roughly 207.9 MB versus 2.25 MB allocations/op. Step median eager/lazy ratio was 1.0067 and StepNLast pp16 1.0115 with unchanged allocation counts; full StepN GPT/Llama parity passed. These are attributed owner measurements, not universal throughput claims or a ranking.", + }, + Analyzer: &analysis.Analyzer{Name: "PS6136", Doc: "source-bound maximum-context output retention with reviewed one-row/bulk policy", Run: runPS6136}, +}) + +func runPS6136(pass *analysis.Pass) (any, error) { + return runPS6136WithContracts(pass, config.Current().OutputWorkspaceContracts) +} + +func runPS6136WithContracts(pass *analysis.Pass, contracts []config.OutputWorkspaceContract) (any, error) { + if len(contracts) == 0 { + return nil, nil + } + context := ps6136ContractsContext(pass) + counts := make(map[string]int) + for index := range contracts { + contract := &contracts[index] + counts[contract.OwnerType+"."+contract.WorkspaceField]++ + } + var pkg *ssa.Package + // Token positions are sparse offsets across unrelated files, not dense IDs. + reported := make(map[token.Pos]bool) + for index := range contracts { + contract := &contracts[index] + if !contract.Valid() || counts[contract.OwnerType+"."+contract.WorkspaceField] != 1 { + continue + } + owner := context.named(contract.OwnerType) + if owner == nil || owner.Obj().Pkg() != pass.Pkg { + continue + } + if pkg == nil { + pkg = ps6136SourcePackage(pass) + } + for _, identity := range contract.ConstructorEntries { + entry := context.callable(identity) + selection := context.selection(pkg, entry, contract) + if selection == nil { + continue + } + proof := selection.constructorAllocation(16384) + if proof == nil || !selection.constructorLifetime(proof, 16384) || !context.nativeBindings(selection, proof, contract) { + continue + } + var consumers []*ps6136ConsumerProof + valid := true + for _, identity := range append(slices.Clone(contract.CommonMethods), contract.BulkMethod) { + method := context.callable(identity) + if method == nil { + valid = false + break + } + function := pkg.Prog.FuncValue(method) + if function == nil { + valid = false + break + } + root := ps6125NewSSAContext(function, nil, nil, 16384) + consumer := selection.consumerCalls(root, contract.Leaves, identity == contract.BulkMethod, func(current *ps6125SSAContext, call *ssa.Call, leaf *config.OutputWorkspaceLeaf) bool { + return context.capacityPrefix(current, call, leaf, selection.width) + }) + if consumer == nil { + valid = false + break + } + consumers = append(consumers, consumer) + } + if !valid || !selection.closedSource(pass, pkg, consumers) || reported[proof.allocation.Pos()] { //perfscan:ignore PS3003 + continue + } + message := "reviewed common one-row output policy retains 4*maximumRows*width bytes, needs 4*width logical bytes, leaves 4*(maximumRows-1)*width idle bytes, and amplifies residency by maximumRows; rare " + contract.BulkMethod + " needs actualRows*width; capacity-wide physical transfers are separate. Consider a private checked bulk generation; preserve allocation/errors, native ranges, completion, publication and release (no automatic fix)" + if contract.ProfileMaximumRows > 0 { + retained, common, idle, amplification, known := proof.residency.evaluate(contract.ProfileMaximumRows, contract.ProfileWidth, math.MaxInt64) + if !known { + continue + } + message = "reviewed profile retains " + strconv.FormatInt(retained, 10) + " output bytes, common logical one-row needs " + strconv.FormatInt(common, 10) + " bytes, idle " + strconv.FormatInt(idle, 10) + " bytes, residency amplification " + strconv.FormatInt(amplification, 10) + "x; " + message + } + message = "conditional on valid checked shapes with maximumRows>1 and width>0: " + message + pass.Report(analysis.Diagnostic{Pos: proof.allocation.Pos(), Message: message}) + reported[proof.allocation.Pos()] = true + } + } + return nil, nil +} + +func (context *ps6136ContractContext) capacityPrefix(current *ps6125SSAContext, call *ssa.Call, leaf *config.OutputWorkspaceLeaf, width *types.Var) bool { + storage, f32 := context.callable(leaf.HostStorageMethod), context.callable(leaf.HostFloat32Method) + if storage == nil || f32 == nil || current == nil || len(current.flow.function.Params) == 0 || current.flow.function.Syntax() == nil { + return false + } + var transfer *ast.CallExpr + ast.Inspect(current.flow.function.Syntax(), func(node ast.Node) bool { + candidate, ok := node.(*ast.CallExpr) + if !ok || call.Pos() < candidate.Pos() || candidate.End() <= call.Pos() { + return true + } + selector, ok := candidate.Fun.(*ast.SelectorExpr) + if ok { + if selection := context.pass.TypesInfo.Selections[selector]; selection != nil && selection.Obj() == call.Call.Method { + transfer = candidate + } + } + return true + }) + return transfer != nil && ps6136HostPrefix(context.pass, transfer, current.flow.function.Params[0].Object(), width, storage, f32) +} diff --git a/checks/ps6136_analyzer_test.go b/checks/ps6136_analyzer_test.go new file mode 100644 index 00000000..7572b4aa --- /dev/null +++ b/checks/ps6136_analyzer_test.go @@ -0,0 +1,112 @@ +package checks + +import ( + "github.com/jxsl13/perfscan/config" + "golang.org/x/tools/go/analysis" + "golang.org/x/tools/go/ssa" + "strings" + "testing" +) + +func ps6136CompleteOwnerContract(pkg *ssa.Package, owner string) config.OutputWorkspaceContract { + c := ps6136NativeOwnerContract(pkg, owner) + local := "github.com/jxsl13/goai/llamagpu." + constructor, entry, head, model := "newGPTDecoder", "NewGPT", "Head", "GPT" + c.ProjectorField = "head" + c.AllocationFunction = local + constructor + if owner == "Decoder" { + constructor, entry, head, model = "newDecoder", "New", "Out", "Llama" + c.ProjectorField = "out" + c.AllocationFunction = local + "Decoder.allocScratch" + } + c.FactoryFunction = local + constructor + c.ConstructorEntries = []string{local + entry} + c.ModelType = "github.com/jxsl13/goai/nlp." + model + c.ModelConfigField = "Config" + c.ModelHeadField = head + c.ConfigMaximumRowsField = "Ctx" + c.ConfigWidthField = "Vocab" + c.HeadShapeMethod = "github.com/jxsl13/goai/tensor.Tensor.Shape" + c.WorkspaceField = "logits" + c.SlotBufferField = "b" + c.MaximumRowsField = "maxLen" + c.WidthField = "v" + c.RetainedListField = "all" + c.BackendOpsField = "ops" + c.BackendAllocatorField = "newBuffer" + c.BackendAllocators = []string{"github.com/jxsl13/goai/backend/metal.NewDeviceBufferF32"} + c.ReleaseMethod = local + owner + ".Release" + c.RetainedBufferReleaseMethod = local + "buffer.Release" + c.NativeBufferReleaseMethod = "github.com/jxsl13/goai/backend/metal.DeviceBuffer.Release" + c.NativeLifetimeSemanticsReviewed = true + c.DominantOneRowPolicyReviewed = true + c.ExplicitRareBulkPolicyReviewed = true + c.FreshFloat32OwnershipReviewed = true + c.ModelProjectorWidthReviewed = true + c.SequentialLifecycleReviewed = true + c.CheckedShapeArithmeticReviewed = true + c.AllProviderBuildPathsReviewed = true + c.ExternalObservationsReviewed = true + c.NativeAccessSemanticsReviewed = true + for index := range c.Leaves { + if c.Leaves[index].Kind == "capacity-transfer" { + c.Leaves[index].HostStorageMethod = "github.com/jxsl13/goai/tensor.Tensor.Storage" + c.Leaves[index].HostFloat32Method = "github.com/jxsl13/goai/tensor.Storage.F32" + } + } + return c +} + +func TestPS6136RegisteredAuthenticOwnerBeforeAfter(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + for _, owner := range []string{"GPTDecoder", "Decoder"} { + t.Run(revision+"/"+owner, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, revision, true, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + c := ps6136CompleteOwnerContract(pkg, owner) + if !c.Valid() { + t.Fatal("authentic selected owner contract invalid") + } + if owner == "GPTDecoder" { + c.ProfileMaximumRows = 1024 + c.ProfileWidth = 50257 + } + var diagnostics []analysis.Diagnostic + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg, Report: func(d analysis.Diagnostic) { diagnostics = append(diagnostics, d) }} + if _, err := runPS6136WithContracts(pass, []config.OutputWorkspaceContract{c}); err != nil { + t.Fatal(err) + } + want := 0 + if revision == "before" { + want = 1 + } + if len(diagnostics) != want { + t.Fatalf("registered owner findings %d want %d", len(diagnostics), want) + } + if want > 0 { + if !strings.Contains(diagnostics[0].Message, "rare "+c.BulkMethod) || !strings.Contains(diagnostics[0].Message, "capacity-wide physical transfers are separate") { + t.Fatal("bulk/physical policy missing") + } + if owner == "GPTDecoder" && !strings.Contains(diagnostics[0].Message, "205852672 output bytes") { + t.Fatal("actual profile retained-byte evidence missing") + } + } + }) + } + } + if !PS6136.NeedsConfig || PS6136.AutoFix { + t.Fatal("unsafe autonomous registration") + } +} + +func TestPS6136UnconfiguredSkipsSSA(t *testing.T) { + t.Parallel() + if _, err := runPS6136WithContracts(&analysis.Pass{}, nil); err != nil { + t.Fatal(err) + } +} diff --git a/checks/ps6136_backend.go b/checks/ps6136_backend.go new file mode 100644 index 00000000..e2803af5 --- /dev/null +++ b/checks/ps6136_backend.go @@ -0,0 +1,248 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// Join the allocator's proved exact native result to the concrete wrapper +// unwrapped by the configured bridge. Equal native-field types in different +// wrapper types cannot satisfy this identity check. +func ps6136BoundBufferWrapper(context *ps6125SSAContext, backend *ssa.Call, wrapper *types.Named, field *types.Var) bool { + if context == nil || backend == nil || wrapper == nil || field == nil { + return false + } + structure, ok := wrapper.Underlying().(*types.Struct) + if !ok || structure.NumFields() != 1 || structure.Field(0) != field { + return false + } + callback := context.call(backend) + if callback == nil { + return false + } + populated := 0 + for _, block := range callback.flow.function.Blocks { + if !callback.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + returned, ok := instruction.(*ssa.Return) + if !ok { + continue + } + if len(returned.Results) != 2 { + return false + } + if constant, ok := returned.Results[0].(*ssa.Const); ok && constant.IsNil() { + continue + } + conversion, ok := returned.Results[0].(*ssa.MakeInterface) + if !ok || !types.Identical(conversion.X.Type(), wrapper) { + return false + } + reference := callback.structField(conversion.X, 0) + result, ok := reference.value.(*ssa.Extract) + if !ok || reference.context != callback || result.Index != 0 { + return false + } + native, ok := result.Tuple.(*ssa.Call) + if !ok || native.Call.StaticCallee() == nil || native.Call.Signature().Results().Len() != 2 || !types.Identical(native.Call.Signature().Results().At(0).Type(), field.Type()) { + return false + } + populated++ + } + } + return populated == 1 +} + +// ps6136BoundAllocator requires the actual closed backend callback to call an +// exact reviewed allocator with the same host descriptor and return that exact +// result, directly or inside a one-field value wrapper. An allocator name or +// a constructor's opaque backendOps formal does not establish this binding. +func ps6136BoundAllocator(context *ps6125SSAContext, backend *ssa.Call, input ps6125SSAReference, identities map[string]bool) bool { + if context == nil || backend == nil || input.value == nil { + return false + } + callback := context.call(backend) + if callback == nil { + return false + } + var native *ssa.Call + for _, block := range callback.flow.function.Blocks { + if !callback.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + call, ok := instruction.(*ssa.Call) + if !ok { + continue + } + callee := call.Call.StaticCallee() + if callee == nil || callee.Object() == nil || !identities[ps6090FunctionID(callee.Object().(*types.Func))] { + continue + } + if native != nil || call.Call.IsInvoke() || len(call.Call.Args) != 1 || callback.reference(call.Call.Args[0]) != input { + return false + } + signature := callee.Signature + if signature.Recv() != nil || signature.Variadic() || signature.TypeParams().Len() != 0 || signature.Params().Len() != 1 || signature.Results().Len() != 2 || + !types.Identical(signature.Params().At(0).Type(), types.NewSlice(types.Typ[types.Float32])) || + !types.Identical(signature.Results().At(1).Type(), types.Universe.Lookup("error").Type()) { + return false + } + native = call + } + } + if native == nil { + return false + } + populated, empty := 0, 0 + var nativeResult *ssa.Extract + var wrapper ps6125SSAReference + for _, block := range callback.flow.function.Blocks { + if !callback.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + returned, ok := instruction.(*ssa.Return) + if !ok { + continue + } + if len(returned.Results) != 2 { + return false + } + if literal, ok := returned.Results[0].(*ssa.Const); ok && literal.IsNil() { + errorValue, ok := callback.reference(returned.Results[1]).value.(*ssa.Extract) + if !ok || errorValue.Tuple != native || errorValue.Index != 1 { + return false + } + if !ps6136NativeErrorBranch(callback, returned, native, true) { + return false + } + empty++ + continue + } + if errorValue, ok := returned.Results[1].(*ssa.Const); !ok || !errorValue.IsNil() { + return false + } + if !ps6136NativeErrorBranch(callback, returned, native, false) { + return false + } + value := returned.Results[0] + if converted, ok := value.(*ssa.MakeInterface); ok { + value = converted.X + } + reference := callback.reference(value) + if structure, ok := value.Type().Underlying().(*types.Struct); ok { + if structure.NumFields() != 1 { + return false + } + reference = callback.structField(value, 0) + wrapper = callback.reference(value) + } + extracted, ok := reference.value.(*ssa.Extract) + if !ok || reference.context != callback || extracted.Tuple != native || extracted.Index != 0 { + return false + } + nativeResult = extracted + populated++ + } + } + if populated != 1 || empty != 1 || nativeResult == nil { + return false + } + users := nativeResult.Referrers() + if users == nil { + return false + } + for _, user := range *users { + if user.Parent() == callback.flow.function && !callback.flow.blocks[user.Block()] { + continue + } + switch use := user.(type) { + case *ssa.DebugRef: + case *ssa.MakeInterface: + if wrapper.value != nil { + return false + } + convertedUsers := use.Referrers() + if convertedUsers == nil { + return false + } + for _, convertedUser := range *convertedUsers { + if _, debug := convertedUser.(*ssa.DebugRef); debug { + continue + } + returned, ok := convertedUser.(*ssa.Return) + if !ok || returned.Parent() != callback.flow.function || len(returned.Results) != 2 || returned.Results[0] != use { + return false + } + } + case *ssa.Return: + if wrapper.value != nil { + return false + } + case *ssa.Store: + address, ok := use.Addr.(*ssa.FieldAddr) + if !ok || wrapper.value == nil { + return false + } + paths := ps6125AccessPaths{flow: callback.flow} + path := paths.resolve(address) + if !path.known || len(path.access.fields) != 1 { + return false + } + // structField already closes the fresh wrapper cell and dominates + // the actual read; unrelated stores may not masquerade as its init. + load, ok := wrapper.value.(*ssa.UnOp) + if !ok || load.X != address.X { + return false + } + default: + return false + } + } + return true +} + +func ps6136NativeErrorBranch(context *ps6125SSAContext, returned *ssa.Return, native *ssa.Call, wantError bool) bool { + block := returned.Block() + if len(block.Preds) != 1 { + return false + } + predecessor := block.Preds[0] + if len(predecessor.Instrs) == 0 { + return false + } + branch, ok := predecessor.Instrs[len(predecessor.Instrs)-1].(*ssa.If) + if !ok { + return false + } + comparison, ok := branch.Cond.(*ssa.BinOp) + if !ok || (comparison.Op != token.EQL && comparison.Op != token.NEQ) { + return false + } + value, nilValue := comparison.X, comparison.Y + if literal, ok := value.(*ssa.Const); ok && literal.IsNil() { + value, nilValue = nilValue, value + } + literal, ok := nilValue.(*ssa.Const) + if !ok || !literal.IsNil() { + return false + } + extracted, ok := context.reference(value).value.(*ssa.Extract) + if !ok || extracted.Tuple != native || extracted.Index != 1 { + return false + } + trueBranch := comparison.Op == token.NEQ + if !wantError { + trueBranch = !trueBranch + } + index := 1 + if trueBranch { + index = 0 + } + return len(predecessor.Succs) == 2 && predecessor.Succs[index] == block +} diff --git a/checks/ps6136_buffer_origin.go b/checks/ps6136_buffer_origin.go new file mode 100644 index 00000000..92b13c23 --- /dev/null +++ b/checks/ps6136_buffer_origin.go @@ -0,0 +1,50 @@ +package checks + +import ( + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// Exact instance/workspace/slot-buffer identity through source interface +// capability assertions. No slice offsets, returned opaque aliases, loads from +// unrelated owners or boolean assertion results are interpreted as a buffer. +func ps6136BufferOrigin(context *ps6125SSAContext, value ssa.Value, owner ps6125SSAReference, ownerType *types.Named, workspace, slot *types.Var, budget int) bool { + if context == nil || value == nil || owner.value == nil || ownerType == nil || workspace == nil || slot == nil || budget <= 0 { + return false + } + reference := context.reference(value) + if reference.value == nil { + return false + } + if reference.context != context || reference.value != value { + return ps6136BufferOrigin(reference.context, reference.value, owner, ownerType, workspace, slot, budget-1) + } + paths := ps6125AccessPaths{flow: context.flow} + path := paths.resolve(value) + if path.known && len(path.access.fields) == 2 && path.access.fields[0] == workspace && path.access.fields[1] == slot && ps6136AccessOwnerRoot(context, path, ownerType) == owner { + return true + } + switch source := value.(type) { + case *ssa.TypeAssert: + return !source.CommaOk && ps6136BufferOrigin(context, source.X, owner, ownerType, workspace, slot, budget-1) + case *ssa.Extract: + assertion, ok := source.Tuple.(*ssa.TypeAssert) + return ok && assertion.CommaOk && source.Index == 0 && ps6136BufferOrigin(context, assertion.X, owner, ownerType, workspace, slot, budget-1) + case *ssa.ChangeInterface: + return ps6136BufferOrigin(context, source.X, owner, ownerType, workspace, slot, budget-1) + case *ssa.Phi: + found := false + for index, predecessor := range source.Block().Preds { + if !context.flow.edges[ps6125SSAEdge{from: predecessor, to: source.Block()}] { + continue + } + if !ps6136BufferOrigin(context, source.Edges[index], owner, ownerType, workspace, slot, budget-1) { + return false + } + found = true + } + return found + } + return false +} diff --git a/checks/ps6136_compiler_test.go b/checks/ps6136_compiler_test.go new file mode 100644 index 00000000..78a27edc --- /dev/null +++ b/checks/ps6136_compiler_test.go @@ -0,0 +1,623 @@ +package checks + +import ( + "bytes" + "go/ast" + "go/format" + "go/importer" + "go/parser" + "go/token" + "go/types" + "os" + "strings" + "testing" + + "golang.org/x/tools/go/ssa" +) + +// The original complete GPT file is compiled without replacing any owner +// body. Supporting declarations and external APIs are explicit scaffolding; +// this is source analysis, not execution or proof of opaque backend semantics. +func ps6136CompileGPT(t *testing.T, revision string) (*analysisOwnerFixture, error) { + return ps6136CompileGPTBackend(t, revision, false) +} + +func ps6136CompileGPTBackend(t *testing.T, revision string, metalEntry bool) (*analysisOwnerFixture, error) { + return ps6136CompileOwners(t, revision, metalEntry, false) +} + +func ps6136CompileOwners(t *testing.T, revision string, metalEntry, shared bool) (*analysisOwnerFixture, error) { + return ps6136CompileOwnersRewrite(t, revision, metalEntry, shared, nil) +} + +// Rewrite is only for explicitly altered adversaries. Authentic replay calls +// the nil-rewrite entry above and still compiles complete original owner bodies. +func ps6136CompileOwnersRewrite(t *testing.T, revision string, metalEntry, shared bool, rewrite func(*token.FileSet, []*ast.File) []*ast.File) (*analysisOwnerFixture, error) { + t.Helper() + fs := token.NewFileSet() + data, err := os.ReadFile("testdata/ps6136-owner/" + revision + "/gpt.go.txt") + if err != nil { + return nil, err + } + gpt, err := parser.ParseFile(fs, "gpt.go", data, parser.SkipObjectResolution) + if err != nil { + return nil, err + } + decoderSource, err := os.ReadFile("testdata/ps6136-owner/" + revision + "/decoder.go.txt") + if err != nil { + return nil, err + } + decoder, err := parser.ParseFile(fs, "decoder.go", decoderSource, parser.SkipObjectResolution) + if err != nil { + return nil, err + } + support, err := parser.ParseFile(fs, "support.go", `package llamagpu +import "github.com/jxsl13/goai/tensor" +func flat1D(*tensor.Tensor)[]float32{return nil} +func flat2D(*tensor.Tensor)[]float32{return nil} +func embedRow(*tensor.Tensor,int,int)[]float32{return nil} +func(l f32Linear)recordMoE(recorder,buffer,buffer,int,buffer,int,int)error{return nil} +`, parser.SkipObjectResolution) + if err != nil { + return nil, err + } + wantedTypes := map[string]bool{"buffer": true, "qweight": true, "linear": true, "recorder": true, "f32Linear": true, "backendOps": true, "bufSlot": true, "growBuffer": true} + wantedFunctions := map[string]bool{"firstErr": true, "logitsForRows": true} + if shared { + wantedFunctions["newDecoder"] = true + wantedFunctions["newDecoderCommon"] = true + wantedFunctions["recordBarrier"] = true + wantedFunctions["boolToInt"] = true + extra, err := parser.ParseFile(fs, "shared-scaffold.go", `package llamagpu +import("fmt";"math";"github.com/jxsl13/goai/backend";"github.com/jxsl13/goai/nlp";"github.com/jxsl13/goai/nn") +var _=fmt.Errorf;var _=math.Sqrt;var _=backend.RoPEFreqs;var _ *nn.RMSNorm +func(q quantLinear)recordAdd(recorder,buffer,buffer,buffer,int,int)error{return nil} +func(q quantLinear)recordMoE(recorder,buffer,buffer,int,buffer,int,int)error{return nil} +func fastTopKSampler(nlp.TokenSampler,int)(*nlp.Sampler,int){return nil,0} +func fastTopPSampler(nlp.TokenSampler,int)(*nlp.Sampler,int){return nil,0} +func sampleTopKCandidates(*nlp.Sampler,[]int32,[]float32)int{return 0} +`, parser.SkipObjectResolution) + if err != nil { + return nil, err + } + // Owner bodies keep their original import binding environment. + for _, declaration := range extra.Decls { + support.Decls = append(support.Decls, declaration) + } + support.Imports = append(support.Imports, extra.Imports...) + deviceSource, err := os.ReadFile("testdata/ps6136-owner/sampling_fastpath.go.txt") + if err != nil { + return nil, err + } + device, err := parser.ParseFile(fs, "sampling_fastpath.go", deviceSource, parser.SkipObjectResolution) + if err != nil { + return nil, err + } + for _, declaration := range device.Decls { + gen, ok := declaration.(*ast.GenDecl) + if !ok || gen.Tok != token.TYPE { + continue + } + for _, spec := range gen.Specs { + typ := spec.(*ast.TypeSpec) + if typ.Name.Name == "deviceTopKer" || typ.Name.Name == "deviceTopPer" { + support.Decls = append(support.Decls, &ast.GenDecl{Tok: token.TYPE, Specs: []ast.Spec{typ}}) + } + } + } + } + for _, declaration := range decoder.Decls { + switch n := declaration.(type) { + case *ast.GenDecl: + if n.Tok == token.CONST { + support.Decls = append(support.Decls, n) + continue + } + for _, spec := range n.Specs { + if typ, ok := spec.(*ast.TypeSpec); ok && (shared || wantedTypes[typ.Name.Name]) { + support.Decls = append(support.Decls, &ast.GenDecl{Tok: token.TYPE, Specs: []ast.Spec{typ}}) + } + } + case *ast.FuncDecl: + wanted := wantedFunctions[n.Name.Name] + if n.Recv != nil && len(n.Recv.List) == 1 { + receiver := n.Recv.List[0].Type + if ptr, ok := receiver.(*ast.StarExpr); ok { + receiver = ptr.X + } + if typ, ok := receiver.(*ast.Ident); ok { + wanted = wanted || typ.Name == "growBuffer" || typ.Name == "f32Linear" && (n.Name.Name == "record" || n.Name.Name == "recordAdd") + wanted = wanted || shared && (typ.Name == "Decoder" || typ.Name == "quantLinear" && n.Name.Name == "record") + } + } + if wanted { + support.Decls = append(support.Decls, n) + } + } + } + imports := &ps6136FixtureImporter{fallback: importer.Default(), packages: map[string]*types.Package{}, sources: map[string]string{}} + info := &types.Info{Types: map[ast.Expr]types.TypeAndValue{}, Uses: map[*ast.Ident]types.Object{}, Defs: map[*ast.Ident]types.Object{}, Selections: map[*ast.SelectorExpr]*types.Selection{}} + files := []*ast.File{gpt, support} + if shared { + // Complete, unmodified owner files: no alternate constructor or + // allocation/flatten helper is omitted from the coverage replay. + samplingSource, err := os.ReadFile("testdata/ps6136-owner/sampling_fastpath.go.txt") + if err != nil { + return nil, err + } + sampling, err := parser.ParseFile(fs, "sampling_fastpath.go", samplingSource, parser.SkipObjectResolution) + if err != nil { + return nil, err + } + files = []*ast.File{gpt, decoder, sampling} + } + if metalEntry { + data, err := os.ReadFile("testdata/ps6136-owner/llamagpu.go.txt") + if err != nil { + return nil, err + } + adapter, err := parser.ParseFile(fs, "llamagpu.go", data, parser.SkipObjectResolution) + if err != nil { + return nil, err + } + var declarations []ast.Decl + for _, declaration := range adapter.Decls { + switch n := declaration.(type) { + case *ast.GenDecl: + if n.Tok == token.IMPORT { + declarations = append(declarations, n) + continue + } + for _, spec := range n.Specs { + if typ, ok := spec.(*ast.TypeSpec); ok && (typ.Name.Name == "mBuf" || typ.Name.Name == "mRec" || shared && typ.Name.Name == "mProfileRec") { + declarations = append(declarations, &ast.GenDecl{Tok: token.TYPE, Specs: []ast.Spec{typ}}) + } + } + case *ast.FuncDecl: + if n.Name.Name == "NewGPT" || shared && n.Name.Name == "New" { + declarations = append(declarations, n) + } + if shared && (n.Name.Name == "NewQuant" || n.Name.Name == "NewQuantF16KV" || n.Name.Name == "newQuantMetal" || n.Name.Name == "newQuantMetalWithMixedQKV" || n.Name.Name == "concurrentMetalDecodeEligible" || n.Name.Name == "ProfileMetalStep") { + declarations = append(declarations, n) + } + if shared && n.Recv != nil && (n.Name.Name == "Finish" || n.Name.Name == "Wait") { + if receiver, ok := n.Recv.List[0].Type.(*ast.Ident); ok && receiver.Name == "mProfileRec" { + declarations = append(declarations, n) + } + } + if n.Name.Name == "mb" || n.Recv != nil && (n.Name.Name == "MatMul" || n.Name.Name == "AddBias") { + declarations = append(declarations, n) + } + } + } + adapter.Decls = declarations + files = append(files, adapter) + if shared { + // Type-only support for an unselected quantized-provider callback; + // this is not source/execution proof of its native implementation. + quantized, err := parser.ParseFile(fs, "quantized-provider-scaffold.go", `package llamagpu;func metalUploadQWeight([]byte,uint32,int,int)(qweight,error){return nil,nil};func metalGroupQWeights([]qweight)(qweight,error){return nil,nil}`, parser.SkipObjectResolution) + if err != nil { + return nil, err + } + files = append(files, quantized) + } + for _, declaration := range decoder.Decls { + gen, ok := declaration.(*ast.GenDecl) + if !ok { + continue + } + for _, spec := range gen.Specs { + typ, ok := spec.(*ast.TypeSpec) + if !ok || typ.Name.Name != "recorder" { + continue + } + api := typ.Type.(*ast.InterfaceType) + for _, method := range api.Methods.List { + if method.Names[0].Name == "MatMul" || method.Names[0].Name == "AddBias" { + continue // The complete original forwarding bodies are above. + } + var signature bytes.Buffer + if err := format.Node(&signature, fs, method.Type); err != nil { + return nil, err + } + body := ps6136ZeroReturnSource(fs, method.Type.(*ast.FuncType)) + generated, err := parser.ParseFile(fs, "recorder-scaffold.go", "package llamagpu;func(mRec) "+method.Names[0].Name+strings.TrimPrefix(signature.String(), "func")+"{"+body+"}", parser.SkipObjectResolution) + if err != nil { + return nil, err + } + files = append(files, generated) + } + } + } + } + if rewrite != nil { + files = rewrite(fs, files) + } + pkg, err := (&types.Config{Importer: imports}).Check("github.com/jxsl13/goai/llamagpu", fs, files, info) + return &analysisOwnerFixture{fs, files, info, pkg, imports.sources}, err +} + +type analysisOwnerFixture struct { + fileset *token.FileSet + files []*ast.File + info *types.Info + pkg *types.Package + sources map[string]string +} + +type ps6136FixtureImporter struct { + fallback types.Importer + packages map[string]*types.Package + sources map[string]string +} + +func (i *ps6136FixtureImporter) Import(path string) (*types.Package, error) { + if pkg := i.packages[path]; pkg != nil { + return pkg, nil + } + var source string + switch path { + case "github.com/jxsl13/goai/tensor": + source = `package tensor +type Tensor struct{};type Shape []int;type Dtype uint8;const F32 Dtype=1 +type Storage struct{};func(*Storage)F32()[]float32{return nil} +func(*Tensor)Shape()Shape{return nil};func(*Tensor)Storage()*Storage{return nil} +func(*Tensor)AtF64(...int)float64{return 0} +func(*Tensor)Slice(int,int,int)(*Tensor,error){return nil,nil} +func(*Tensor)Dtype()Dtype{return F32};func(*Tensor)IsContiguous()bool{return true} +func(*Tensor)Offset()int{return 0};func(t *Tensor)Cast(Dtype)*Tensor{return t} +func New(Dtype,Shape)*Tensor{return nil} +func(*Tensor)Transpose(int,int)(*Tensor,error){return nil,nil} +` + case "github.com/jxsl13/goai/nlp": + api, err := os.ReadFile("testdata/ps6136-owner/model-api.go.txt") + if err != nil { + return nil, err + } + source = string(api) + ` +type TokenSampler interface{SampleWithHistory([]float64,[]int)int} +type Sampler struct{Temperature,TopP,MinP,Epsilon,Eta,Typical,RepeatPenalty,FreqPenalty,PresencePenalty,DRYMultiplier,XTCProbability float64;TopK,TopNSigma int} +func(*Sampler)Sample([]float64)int{return 0} +func(*Sampler)SampleWithHistory([]float64,[]int)int{return 0} +func(*Sampler)SampleTopPFromCandidates([]float64,[]int32,float64,float64)(int,bool){return 0,false} +` + case "github.com/jxsl13/goai/backend": + source = `package backend;type Kind string;type Name string;type Op int;type ResidentWeight interface{};const Metal Kind="metal";var ErrQuantUnsupported error +type RoPEAttrs struct{Base float64;Heads int} +func RoPEFreqs(int,RoPEAttrs)([]float64,float64){return nil,0} +func ALiBiSlopes(int)[]float64{return nil} +` + case "github.com/jxsl13/goai/nn": + api, err := os.ReadFile("testdata/ps6136-owner/nn-api.go.txt") + if err != nil { + return nil, err + } + source = string(api) + case "github.com/jxsl13/goai/format/gguf": + source = `package gguf;type QuantType uint32` + case "github.com/jxsl13/goai/backend/metal": + source = `package metal +type DeviceBuffer struct{} +func(*DeviceBuffer)UploadF32([]float32)error{return nil} +func(*DeviceBuffer)DownloadF32([]float32)error{return nil} +func(*DeviceBuffer)Release(){} +func(*DeviceBuffer)Len()int{return 0};func(*DeviceBuffer)ByteLen()int{return 0} +func(*DeviceBuffer)TopKN(int,int)([]int32,[]float32,error){return nil,nil,nil} +type Recorder struct{} +type RecorderProfile struct{} +func(*Recorder)MatMul(*DeviceBuffer,*DeviceBuffer,*DeviceBuffer,int,int,int)error{return nil} +func(*Recorder)AddBias(*DeviceBuffer,*DeviceBuffer,*DeviceBuffer,int,int)error{return nil} +func(*Recorder)Finish()error{return nil};func(*Recorder)Wait()error{return nil} +func(*Recorder)Profile()(RecorderProfile,error){return RecorderProfile{},nil} +func Available()bool{return true} +func NewDeviceBufferF32([]float32)(*DeviceBuffer,error){return nil,nil} +func NewDeviceBufferF16Zeros(int)(*DeviceBuffer,error){return nil,nil} +func NewRecorder()(*Recorder,error){return nil,nil} +func NewConcurrentRecorder()(*Recorder,error){return nil,nil} +func NewProfilingRecorder(int)(*Recorder,error){return nil,nil} +` + default: + return i.fallback.Import(path) + } + fs := token.NewFileSet() + file, err := parser.ParseFile(fs, path+".go", source, parser.SkipObjectResolution) + if err != nil { + return nil, err + } + pkg, err := (&types.Config{Importer: i}).Check(path, fs, []*ast.File{file}, nil) + if err == nil { + i.packages[path] = pkg + if i.sources != nil { + i.sources[path] = source + } + } + return pkg, err +} + +func ps6136ZeroReturnSource(fs *token.FileSet, signature *ast.FuncType) string { + if signature.Results == nil { + return "" + } + var source strings.Builder + var names []string + for _, field := range signature.Results.List { + count := len(field.Names) + if count == 0 { + count = 1 + } + var typ bytes.Buffer + if err := format.Node(&typ, fs, field.Type); err != nil { + panic(err) + } + for range count { + name := "result" + string(rune('A'+len(names))) + source.WriteString("var " + name + " " + typ.String() + ";") + names = append(names, name) + } + } + source.WriteString("return " + strings.Join(names, ",")) + return source.String() +} + +func TestPS6136CompleteGPTCompilerFixture(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + t.Run(revision, func(t *testing.T) { + t.Parallel() + if _, err := ps6136CompileGPT(t, revision); err != nil { + t.Fatal(err) + } + }) + } +} + +func TestPS6136CompleteSharedDecoderCompilerFixture(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + t.Run(revision, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, revision, false, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + if pkg.Func("newDecoder") == nil { + t.Fatal("actual shared constructor absent") + } + for _, name := range []string{"newDecoderCommon", "newStableLMDecoder", "newCohereDecoder", "newNemotronDecoder", "newGemmaDecoder", "newDeepSeekV2Decoder", "newMambaDecoder", "newJambaDecoder", "newMamba2Decoder", "newRWKVDecoder", "newFalconDecoder", "newOLMo2Decoder", "newGemma2Decoder", "newMPTDecoder", "newMixtralDecoder", "newQwen2MoEDecoder", "newGraniteMoEDecoder", "newOLMoEDecoder", "newStarCoder2Decoder", "newPhiDecoder", "newGPTNeoXDecoder", "newQuantDecoder"} { + if pkg.Func(name) == nil { + t.Fatalf("actual alternate constructor %s absent", name) + } + } + owner := pkg.Pkg.Scope().Lookup("Decoder").Type().(*types.Named) + for _, name := range []string{"allocScratch", "allocMambaScratch", "allocMamba2Scratch", "allocRWKVScratch", "mkBuf", "recordLogits", "encodeStep", "stepInto", "stepN", "Step", "StepN", "StepNLast", "Generate", "Release"} { + if pkg.Prog.LookupMethod(types.NewPointer(owner), pkg.Pkg, name) == nil { + t.Fatalf("actual owner method %s absent", name) + } + } + }) + } +} + +func ps6136FixtureSSA(fixture *analysisOwnerFixture) *ssa.Package { + program := ssa.NewProgram(fixture.fileset, ssa.SanityCheckFunctions) + seen := map[*types.Package]bool{} + var create func(*types.Package) + create = func(pkg *types.Package) { + if seen[pkg] { + return + } + seen[pkg] = true + for _, imported := range pkg.Imports() { + create(imported) + } + if pkg != fixture.pkg { + program.CreatePackage(pkg, nil, nil, true) + } + } + create(fixture.pkg) + converted := program.CreatePackage(fixture.pkg, fixture.files, fixture.info, true) + converted.Build() + return converted +} + +func TestPS6136AuthenticGPTConstructorFactory(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + t.Run(revision, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileGPT(t, revision) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + constructor := pkg.Func("newGPTDecoder") + root := ps6125NewSSAContext(constructor, nil, nil, 32) + owner := pkg.Pkg.Scope().Lookup("GPTDecoder").Type().(*types.Named) + fieldObject, _, _ := types.LookupFieldOrMethod(owner, true, pkg.Pkg, "logits") + workspace := fieldObject.(*types.Var) + listObject, _, _ := types.LookupFieldOrMethod(owner, true, pkg.Pkg, "all") + list := listObject.(*types.Var) + allocatorObject, _, _ := types.LookupFieldOrMethod(pkg.Pkg.Scope().Lookup("backendOps").Type(), true, pkg.Pkg, "newBuffer") + allocator := allocatorObject.(*types.Var) + slot := pkg.Pkg.Scope().Lookup("bufSlot").Type().Underlying().(*types.Struct).Field(0) + paths := ps6125AccessPaths{flow: root.flow} + found := 0 + for _, block := range constructor.Blocks { + for _, instruction := range block.Instrs { + store, ok := instruction.(*ssa.Store) + if !ok || !root.flow.blocks[block] { + continue + } + path := paths.resolve(store.Addr) + if !path.known || len(path.access.fields) != 1 || path.access.fields[0] != workspace { + continue + } + call, ok := store.Val.(*ssa.Call) + if !ok || len(call.Call.Args) != 1 { + t.Fatal("actual workspace store lost factory input") + } + factory := root.call(call) + if factory == nil { + t.Fatal("actual captured factory context unavailable") + } + backend := ps6136FactoryResult(factory, root.reference(call.Call.Args[0]), allocator, slot) + if backend == nil { + t.Fatal("actual backend input and returned slot provenance unproved") + } + if cell := ps6136FactoryErrorCell(factory, backend); cell.value == nil || cell.context != root { + t.Fatal("actual allocator failure does not feed the same constructor error cell") + } + identity := ps6136AccessOwnerRoot(root, path, owner) + method, _, _ := types.LookupFieldOrMethod(owner, true, pkg.Pkg, "Release") + if !ps6136ConstructorErrorBarrier(root, identity, owner, ps6136FactoryErrorCell(factory, backend), ps6090FunctionID(method.(*types.Func))) { + t.Fatal("actual same-owner final failure release/publication barrier unproved") + } + if !ps6136FactoryRetention(factory, backend, identity, owner, list, slot) { + t.Fatal("actual source result is not exclusively appended to same-owner retention and returned slot") + } + found++ + } + } + if found != 1 { + t.Fatalf("actual constructor workspace factory count %d", found) + } + }) + } +} + +func TestPS6136AuthenticPublicGPTBackendBinding(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + t.Run(revision, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileGPTBackend(t, revision, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + entry := ps6125NewSSAContext(pkg.Func("NewGPT"), nil, nil, 64) + constructor := pkg.Func("newGPTDecoder") + var root *ps6125SSAContext + for _, call := range ps6125ContextCalls(entry) { + if call.Call.StaticCallee() == constructor { + root = entry.call(call) + } + } + if root == nil { + t.Fatal("actual public constructor entry binding unavailable") + } + owner := pkg.Pkg.Scope().Lookup("GPTDecoder").Type().(*types.Named) + workspaceObject, _, _ := types.LookupFieldOrMethod(owner, true, pkg.Pkg, "logits") + allocatorObject, _, _ := types.LookupFieldOrMethod(pkg.Pkg.Scope().Lookup("backendOps").Type(), true, pkg.Pkg, "newBuffer") + slot := pkg.Pkg.Scope().Lookup("bufSlot").Type().Underlying().(*types.Struct).Field(0) + paths := ps6125AccessPaths{flow: root.flow} + found := 0 + for _, block := range constructor.Blocks { + for _, instruction := range block.Instrs { + store, ok := instruction.(*ssa.Store) + if !ok || !root.flow.blocks[block] { + continue + } + path := paths.resolve(store.Addr) + if !path.known || len(path.access.fields) != 1 || path.access.fields[0] != workspaceObject { + continue + } + call := store.Val.(*ssa.Call) + input := root.reference(call.Call.Args[0]) + factory := root.call(call) + backend := ps6136FactoryResult(factory, input, allocatorObject.(*types.Var), slot) + if backend == nil || !ps6136BoundAllocator(factory, backend, input, map[string]bool{"github.com/jxsl13/goai/backend/metal.NewDeviceBufferF32": true}) { + t.Fatal("actual public entry -> allocation callback -> native allocator descriptor/result binding unproved") + } + if ps6136BoundAllocator(factory, backend, input, map[string]bool{"other/backend.NewDeviceBufferF32": true}) { + t.Fatal("wrong exact backend identity accepted") + } + found++ + } + } + if found != 1 { + t.Fatalf("public constructor allocation count %d", found) + } + }) + } +} + +func TestPS6136AuthenticSharedConstructorFactory(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + t.Run(revision, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, revision, false, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + owner := pkg.Pkg.Scope().Lookup("Decoder").Type().(*types.Named) + workspace, _, _ := types.LookupFieldOrMethod(owner, true, pkg.Pkg, "logits") + list, _, _ := types.LookupFieldOrMethod(owner, true, pkg.Pkg, "all") + allocator, _, _ := types.LookupFieldOrMethod(pkg.Pkg.Scope().Lookup("backendOps").Type(), true, pkg.Pkg, "newBuffer") + slot := pkg.Pkg.Scope().Lookup("bufSlot").Type().Underlying().(*types.Struct).Field(0) + root := ps6125NewSSAContext(pkg.Func("newDecoder"), nil, nil, 96) + found := 0 + var visit func(*ps6125SSAContext) + visit = func(context *ps6125SSAContext) { + paths := ps6125AccessPaths{flow: context.flow} + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + store, ok := instruction.(*ssa.Store) + if !ok { + continue + } + path := paths.resolve(store.Addr) + if !path.known || len(path.access.fields) != 1 || path.access.fields[0] != workspace { + continue + } + call, ok := store.Val.(*ssa.Call) + if !ok || len(call.Call.Args) != 1 { + t.Fatal("shared allocation lost actual factory input") + } + factory := ps6136Call(context, call) + if factory == nil { + t.Fatalf("shared mkBuf returned closure context unavailable: value %T %s; reference %T %v", call.Call.Value, call.Call.Value, context.reference(call.Call.Value).value, context.reference(call.Call.Value).value) + } + backend := ps6136FactoryResult(factory, context.reference(call.Call.Args[0]), allocator.(*types.Var), slot) + identity := ps6136AccessOwnerRoot(context, path, owner) + if backend == nil || !ps6136FactoryRetention(factory, backend, identity, owner, list.(*types.Var), slot) { + factoryPaths := ps6125AccessPaths{flow: factory.flow} + for _, factoryBlock := range factory.flow.function.Blocks { + for _, factoryInstruction := range factoryBlock.Instrs { + if write, ok := factoryInstruction.(*ssa.Store); ok { + p := factoryPaths.resolve(write.Addr) + if p.known && len(p.access.fields) == 1 && p.access.fields[0] == list { + t.Logf("retention root %v path %v", ps6136AccessOwnerRoot(factory, p, owner), p) + } + } + } + } + t.Fatalf("actual shared factory input/result and same-owner retention unproved: backend %v identity %v", backend, identity) + } + found++ + } + } + for _, call := range ps6125ContextCalls(context) { + callee := call.Call.StaticCallee() + if callee != nil && callee.Name() == "allocScratch" { + child := context.call(call) + if child == nil { + t.Fatal("shared constructor allocScratch context unavailable") + } + visit(child) + } + } + } + visit(root) + if found != 1 { + t.Fatalf("shared constructor workspace allocations %d", found) + } + }) + } +} diff --git a/checks/ps6136_compound_review_test.go b/checks/ps6136_compound_review_test.go new file mode 100644 index 00000000..8b205764 --- /dev/null +++ b/checks/ps6136_compound_review_test.go @@ -0,0 +1,107 @@ +package checks + +import ( + "go/ast" + "go/parser" + "go/token" + "testing" + + "github.com/jxsl13/perfscan/config" + "golang.org/x/tools/go/analysis" +) + +func TestPS6136RegisteredCompoundReviewGuards(t *testing.T) { + t.Parallel() + for _, test := range []struct{ owner, mutation string }{ + {"Decoder", "destination_index_escape"}, + {"Decoder", "unknown_bias_recorder"}, + {"GPTDecoder", "unknown_projection_recorder"}, + {"Decoder", "unknown_projection_recorder"}, + {"GPTDecoder", "no_actual_bulk_extent"}, + {"Decoder", "no_actual_bulk_extent"}, + } { + t.Run(test.owner+"/"+test.mutation, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwnersRewrite(t, "before", true, true, func(fs *token.FileSet, files []*ast.File) []*ast.File { + changed := 0 + for _, file := range files { + for _, declaration := range file.Decls { + function, ok := declaration.(*ast.FuncDecl) + if !ok || function.Recv == nil { + continue + } + pointer, ok := function.Recv.List[0].Type.(*ast.StarExpr) + if !ok { + continue + } + receiver, ok := pointer.X.(*ast.Ident) + if !ok || receiver.Name != test.owner { + continue + } + ast.Inspect(function.Body, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if ok { + selector, selected := call.Fun.(*ast.SelectorExpr) + if selected && test.mutation == "unknown_bias_recorder" && function.Name.Name == "recordLogits" && selector.Sel.Name == "AddBias" { + selector.X = &ast.CallExpr{Fun: &ast.Ident{Name: "unknownRecorder"}} + changed++ + } + projectionFunction := function.Name.Name == "recordLogits" || test.owner == "GPTDecoder" && function.Name.Name == "Step" + var projector *ast.SelectorExpr + if selected { + projector, _ = selector.X.(*ast.SelectorExpr) + } + if selected && test.mutation == "unknown_projection_recorder" && projectionFunction && selector.Sel.Name == "record" && projector != nil && (projector.Sel.Name == "head" || projector.Sel.Name == "out") { + call.Args[0] = &ast.CallExpr{Fun: &ast.Ident{Name: "unknownRecorder"}} + changed++ + } + if selected && test.mutation == "no_actual_bulk_extent" && function.Name.Name == "StepN" && len(call.Args) == 3 { + call.Args[2] = &ast.Ident{Name: "true"} + changed++ + } + } + if test.mutation == "destination_index_escape" && function.Name.Name == "Generate" { + assignment, ok := node.(*ast.AssignStmt) + if ok && len(assignment.Lhs) == 1 && len(assignment.Rhs) == 1 { + conversion, ok := assignment.Rhs[0].(*ast.CallExpr) + if ok && len(conversion.Args) == 1 { + index, ok := conversion.Args[0].(*ast.IndexExpr) + if ok { + if source, ok := index.X.(*ast.Ident); ok && source.Name == "lf" { + destination := assignment.Lhs[0].(*ast.IndexExpr) + iterator := destination.Index.(*ast.Ident) + destination.X = &ast.CallExpr{Fun: &ast.Ident{Name: "unknownDestination"}, Args: []ast.Expr{&ast.UnaryExpr{Op: token.AND, X: &ast.Ident{Name: iterator.Name}}}} + changed++ + } + } + } + } + } + return true + }) + } + } + if changed != 1 { + t.Fatalf("compound mutation changes %d want1", changed) + } + extra, err := parser.ParseFile(fs, "review-helper.go", `package llamagpu;var externalRecorder recorder;func unknownRecorder()recorder{return externalRecorder};func unknownDestination(p *int)[]float64{*p=1<<20;return make([]float64,(1<<20)+1)}`, parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + return append(files, extra) + }) + if err != nil { + t.Fatal(err) + } + c := ps6136CompleteOwnerContract(ps6136FixtureSSA(fixture), test.owner) + var diagnostics []analysis.Diagnostic + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg, Report: func(d analysis.Diagnostic) { diagnostics = append(diagnostics, d) }} + if _, err := runPS6136WithContracts(pass, []config.OutputWorkspaceContract{c}); err != nil { + t.Fatal(err) + } + if len(diagnostics) != 0 { + t.Fatalf("unproved compound observation produced %d findings", len(diagnostics)) + } + }) + } +} diff --git a/checks/ps6136_constructor_writes.go b/checks/ps6136_constructor_writes.go new file mode 100644 index 00000000..d295bd0a --- /dev/null +++ b/checks/ps6136_constructor_writes.go @@ -0,0 +1,136 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136ConstructorWrites is a fresh-owner source census, not permission to +// assume selected constructor geometry. Every initialization target is tied +// to an actual invocation's fresh allocation. Selected final field values and +// post-publication invariance must be checked independently. +type ps6136ConstructorWrites struct { + positions map[token.Pos]bool + writers map[*ssa.Function]bool + contexts map[*ssa.Function]bool + calls map[*ssa.Call]bool + reads map[token.Pos]bool + readFields map[token.Pos]*types.Var +} + +func ps6136FreshConstructorWrites(pkg *ssa.Package, owner *types.Named, fields map[*types.Var]bool) *ps6136ConstructorWrites { + inventory := ps6136ConstructorInventory(pkg, owner) + if inventory == nil { + return nil + } + result := &ps6136ConstructorWrites{positions: make(map[token.Pos]bool), writers: make(map[*ssa.Function]bool), contexts: make(map[*ssa.Function]bool), calls: make(map[*ssa.Call]bool), reads: make(map[token.Pos]bool), readFields: make(map[token.Pos]*types.Var)} + remaining := 16384 + for function := range inventory { + pending := []*ps6125SSAContext{ps6125NewSSAContext(function, nil, nil, 1024)} + seen := make(map[*ps6125SSAContext]bool) + for len(pending) != 0 { + context := pending[len(pending)-1] + pending = pending[:len(pending)-1] + if seen[context] { + continue + } + seen[context] = true + remaining-- + if remaining < 0 { + return nil + } + result.contexts[context.flow.function] = true + paths := ps6125AccessPaths{flow: context.flow} + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + if load, ok := instruction.(*ssa.UnOp); ok && load.Op == token.MUL { + path := paths.resolve(load.X) + if path.known && len(path.access.fields) == 1 && fields[path.access.fields[0]] { + if _, fresh := ps6136AccessOwnerRoot(context, path, owner).value.(*ssa.Alloc); !fresh { + return nil + } + result.reads[load.Pos()] = true + result.readFields[load.Pos()] = path.access.fields[0] + } + } + if write, ok := instruction.(*ssa.Store); ok { + if address, ok := write.Addr.(*ssa.FieldAddr); ok { + pointer, ok := address.X.Type().Underlying().(*types.Pointer) + if ok { + if structure, ok := pointer.Elem().Underlying().(*types.Struct); ok && fields[structure.Field(address.Field)] { + root := address.X + for { + parent, ok := root.(*ssa.FieldAddr) + if !ok { + break + } + root = parent.X + } + if types.Identical(root.Type(), types.NewPointer(owner)) { + if _, fresh := ps6136OwnerRoot(context, root, owner).value.(*ssa.Alloc); !fresh || !write.Pos().IsValid() { + return nil + } + result.positions[write.Pos()] = true + result.writers[context.flow.function] = true + continue + } + if _, localValue := root.(*ssa.Alloc); localValue && !types.Identical(root.Type(), types.NewPointer(owner)) { + continue + } + } + } + } + if address, ok := write.Addr.(*ssa.FieldAddr); ok && types.Identical(address.X.Type(), types.NewPointer(owner)) { + structure := owner.Underlying().(*types.Struct) + if fields[structure.Field(address.Field)] { + if _, fresh := ps6136OwnerRoot(context, address.X, owner).value.(*ssa.Alloc); !fresh || !write.Pos().IsValid() { + return nil + } + result.positions[write.Pos()] = true + result.writers[context.flow.function] = true + continue + } + } + path := paths.resolve(write.Addr) + if !path.known || len(path.access.fields) == 0 || !fields[path.access.fields[len(path.access.fields)-1]] { + continue + } + identity := ps6136AccessOwnerRoot(context, path, owner) + if _, fresh := identity.value.(*ssa.Alloc); !fresh || !write.Pos().IsValid() { + return nil + } + result.positions[write.Pos()] = true + result.writers[context.flow.function] = true + } + call, ok := instruction.(*ssa.Call) + if !ok { + continue + } + if child := ps6136Call(context, call); child != nil { + for _, argument := range call.Call.Args { + if types.Identical(argument.Type(), types.NewPointer(owner)) { + if _, fresh := ps6136OwnerRoot(context, argument, owner).value.(*ssa.Alloc); !fresh { + return nil + } + result.calls[call] = true + } + } + pending = append(pending, child) + } else { + for _, argument := range call.Call.Args { + if types.Identical(argument.Type(), types.NewPointer(owner)) || types.Identical(argument.Type(), owner) { + return nil // Unknown source ownership edge. + } + } + } + } + } + } + } + return result +} diff --git a/checks/ps6136_constructors.go b/checks/ps6136_constructors.go new file mode 100644 index 00000000..4d62a415 --- /dev/null +++ b/checks/ps6136_constructors.go @@ -0,0 +1,63 @@ +package checks + +import ( + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136ConstructorInventory is unconditional over source-visible owner-returning +// functions. Every non-nil result must originate in a fresh, source-visible +// allocation; aliases, global/opaque returns or ambiguous instances reject the +// entire inventory. This does not establish selected-model width, allocation +// cleanup, provider identity or safe subsequent field writes. +func ps6136ConstructorInventory(pkg *ssa.Package, owner *types.Named) map[*ssa.Function]ps6125SSAReference { + if pkg == nil || owner == nil { + return nil + } + functions := ps6136SourceFunctions(pkg) + result := make(map[*ssa.Function]ps6125SSAReference, len(functions)) + for _, function := range functions { + for index := 1; index < function.Signature.Results().Len(); index++ { + if types.Identical(function.Signature.Results().At(index).Type(), types.NewPointer(owner)) { + return nil // Unsupported alternate owner-result roles are not omitted. + } + } + if function.Signature.Results().Len() == 0 || !types.Identical(function.Signature.Results().At(0).Type(), types.NewPointer(owner)) { + continue + } + if len(function.Blocks) == 0 || function.Signature.Variadic() || function.Signature.TypeParams().Len() != 0 { + return nil + } + context := ps6125NewSSAContext(function, nil, nil, 512) + var identity ps6125SSAReference + for _, block := range function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + returned, ok := instruction.(*ssa.Return) + if !ok || len(returned.Results) == 0 { + continue + } + value := returned.Results[0] + if constant, ok := value.(*ssa.Const); ok && constant.IsNil() { + continue + } + root := ps6136OwnerRoot(context, value, owner) + if _, fresh := root.value.(*ssa.Alloc); !fresh || identity.value != nil && identity != root { + return nil + } + identity = root + } + } + if identity.value == nil { + return nil + } + result[function] = identity + } + if len(result) == 0 { + return nil + } + return result +} diff --git a/checks/ps6136_constructors_test.go b/checks/ps6136_constructors_test.go new file mode 100644 index 00000000..d6777905 --- /dev/null +++ b/checks/ps6136_constructors_test.go @@ -0,0 +1,94 @@ +package checks + +import ( + "go/types" + "testing" + + "golang.org/x/tools/go/analysis" +) + +func TestPS6136AuthenticCompleteConstructorInventory(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + t.Run(revision, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, revision, true, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + for _, test := range []struct { + owner string + count int + }{{"GPTDecoder", 2}, {"Decoder", 28}} { + owner := pkg.Pkg.Scope().Lookup(test.owner).Type().(*types.Named) + inventory := ps6136ConstructorInventory(pkg, owner) + if len(inventory) != test.count { + t.Fatalf("%s full source constructor inventory=%d want %d", test.owner, len(inventory), test.count) + } + } + owner := pkg.Pkg.Scope().Lookup("Decoder").Type().(*types.Named) + profile, _, _ := types.LookupFieldOrMethod(owner, true, pkg.Pkg, "ProfileMetalStep") + if profile == nil || pkg.Prog.FuncValue(profile.(*types.Func)) == nil || len(pkg.Prog.FuncValue(profile.(*types.Func)).Blocks) == 0 { + t.Fatal("actual one-row profiling API body missing from complete owner replay") + } + }) + } +} + +func TestPS6136ConstructorInventoryBoundaries(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name, extra string + valid bool + }{ + {"fresh_only", "", true}, + {"borrowed", `func borrowed(d *owner)*owner{return d}`, false}, + {"global", `var d=&owner{};func existing()*owner{return d}`, false}, + {"opaque", `func opaque()*owner;func unknown()*owner{return opaque()}`, false}, + {"two_instances", `var flag bool;func ambiguous()*owner{if flag{return &owner{}};return &owner{}}`, false}, + {"wrapper_borrowed", `type wrapper struct{d *owner};func(w wrapper)existing()*owner{return w.d}`, false}, + {"anonymous_borrowed", `func getters(d *owner)func()*owner{return func()*owner{return d}}`, false}, + {"alternate_result_role", `func reversed(d *owner)(error,*owner){return nil,d}`, false}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + pkg := ps6125TestSSA(t, `package extent;type owner struct{width int};func fresh()*owner{return &owner{}};`+test.extra) + owner := pkg.Pkg.Scope().Lookup("owner").Type().(*types.Named) + if got := ps6136ConstructorInventory(pkg, owner) != nil; got != test.valid { + t.Fatalf("closed constructor inventory=%v want %v", got, test.valid) + } + }) + } +} + +func TestPS6136AuthenticFreshGeometryWrites(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, "before", true, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + for _, name := range []string{"GPTDecoder", "Decoder"} { + owner := pkg.Pkg.Scope().Lookup(name).Type().(*types.Named) + fields := make(map[*types.Var]bool) + for _, name := range []string{"v", "maxLen", "logits", "ops"} { + field, _, _ := types.LookupFieldOrMethod(owner, true, pkg.Pkg, name) + fields[field.(*types.Var)] = true + } + writes := ps6136FreshConstructorWrites(pkg, owner, fields) + if writes == nil || len(writes.positions) == 0 { + t.Fatal("actual constructor geometry/workspace writes lack closed fresh ownership contexts") + } + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg} + initializers := ps6136ConstructorWriteNodes(pass, pkg, owner, fields, writes) + if initializers == nil { + t.Fatal("actual constructor initialization targets or source caller closure missing") + } + for field := range fields { + if !ps6136FieldEffects(pass, field, initializers) { + t.Fatalf("actual %s.%s has unclassified direct field effects", name, field.Name()) + } + } + } +} diff --git a/checks/ps6136_consumer_walk.go b/checks/ps6136_consumer_walk.go new file mode 100644 index 00000000..ec4fbc6e --- /dev/null +++ b/checks/ps6136_consumer_walk.go @@ -0,0 +1,30 @@ +package checks + +import "golang.org/x/tools/go/ssa" + +// ps6136WalkConsumerCalls retains invocation-specific branches and arguments. +// A visitor may stop at a separately proved leaf. Unresolved calls are delivered +// to the visitor, not treated as harmless: complete observation closure remains +// mandatory before a finding. Exhausting the traversal budget rejects the walk. +func ps6136WalkConsumerCalls(context *ps6125SSAContext, visit func(*ps6125SSAContext, *ssa.Call) bool, remaining *int) bool { + if context == nil || context.flow == nil || visit == nil || remaining == nil || *remaining <= 0 { + return false + } + *remaining-- + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + call, ok := instruction.(*ssa.Call) + if !ok || visit(context, call) { + continue + } + child := ps6136Call(context, call) + if child != nil && !ps6136WalkConsumerCalls(child, visit, remaining) { + return false + } + } + } + return true +} diff --git a/checks/ps6136_consumer_walk_test.go b/checks/ps6136_consumer_walk_test.go new file mode 100644 index 00000000..97d33e5e --- /dev/null +++ b/checks/ps6136_consumer_walk_test.go @@ -0,0 +1,76 @@ +package checks + +import ( + "go/types" + "testing" + + "github.com/jxsl13/perfscan/config" + "golang.org/x/tools/go/ssa" +) + +func TestPS6136AuthenticLastAndBulkConsumerSpecialization(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, "before", true, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + for _, ownerName := range []string{"GPTDecoder", "Decoder"} { + for _, methodName := range []string{"StepNLast", "StepN"} { + t.Run(ownerName+"/"+methodName, func(t *testing.T) { + t.Parallel() + ownerType := pkg.Pkg.Scope().Lookup(ownerName).Type().(*types.Named) + method, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, methodName) + function := pkg.Prog.FuncValue(method.(*types.Func)) + field := func(name string) *types.Var { + object, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, name) + return object.(*types.Var) + } + projectorName := "head" + if ownerName == "Decoder" { + projectorName = "out" + } + workspace, width, projector := field("logits"), field("v"), field(projectorName) + slot := pkg.Pkg.Scope().Lookup("bufSlot").Type().Underlying().(*types.Struct).Field(0) + var pool ps6125ExtentPool + widthExtent := ps6125SymbolicExtent(pool.identity(function.Params[0].Object(), []*types.Var{width}, false)) + tokenExtent := ps6125SymbolicExtent(pool.identity(function.Params[1].Object(), nil, true)) + bulkExtent := ps6125MultiplyExtents(widthExtent, tokenExtent) + context := ps6125NewSSAContext(function, nil, map[*ssa.Parameter]ps6125Extent{function.Params[1]: tokenExtent}, 16384) + facts := &ps6136Extents{owner: context.reference(function.Params[0]), ownerType: ownerType, immutable: map[*types.Var]ps6125Extent{width: widthExtent}, remaining: 8192} + leaves, bulk := 0, 0 + remaining := 8192 + if !ps6136WalkConsumerCalls(context, func(current *ps6125SSAContext, call *ssa.Call) bool { + if !call.Call.IsInvoke() { + return false + } + leaf := config.OutputWorkspaceLeaf{Method: ps6090FunctionID(call.Call.Method), RowsArgument: -1, WidthArgument: -1, HostArgument: -1, SemanticsReviewed: true} + switch call.Call.Method.Name() { + case "record": + leaf.Kind, leaf.BufferArguments, leaf.RowsArgument = "projection", []int{2}, 3 + case "DownloadF32": + leaf.Kind, leaf.BufferArguments, leaf.HostArgument = "download", []int{-1}, 0 + default: + return false + } + got, ok := ps6136LeafExtent(current, call, &leaf, facts, workspace, slot, projector, width) + if !ok { + return false // unrelated layer/projector + } + leaves++ + if ps6125SameExtent(got, bulkExtent) { + bulk++ + } else if !ps6125SameExtent(got, widthExtent) { + t.Fatal("unproved owner leaf extent") + } + return true + }, &remaining) { + t.Fatal("consumer traversal exhausted") + } + if leaves < 2 || methodName == "StepNLast" && bulk != 0 || methodName == "StepN" && bulk < 2 { + t.Fatalf("source leaves=%d bulk=%d", leaves, bulk) + } + }) + } + } +} diff --git a/checks/ps6136_contract_context.go b/checks/ps6136_contract_context.go new file mode 100644 index 00000000..5028e33c --- /dev/null +++ b/checks/ps6136_contract_context.go @@ -0,0 +1,123 @@ +package checks + +import ( + "go/ast" + "go/types" + "strings" + + "golang.org/x/tools/go/analysis" + "golang.org/x/tools/go/ssa" +) + +type ps6136ContractContext struct { + pass *analysis.Pass + packages map[string]*types.Package + declarations map[*types.Func]*ast.FuncDecl +} + +func ps6136ContractsContext(pass *analysis.Pass) *ps6136ContractContext { + context := &ps6136ContractContext{pass: pass, packages: make(map[string]*types.Package), declarations: make(map[*types.Func]*ast.FuncDecl)} + var add func(*types.Package) + add = func(pkg *types.Package) { + if pkg == nil || context.packages[pkg.Path()] != nil { + return + } + context.packages[pkg.Path()] = pkg + for _, imported := range pkg.Imports() { + add(imported) + } + } + add(pass.Pkg) + for _, file := range pass.Files { + for _, declaration := range file.Decls { + if function, ok := declaration.(*ast.FuncDecl); ok { + if object, ok := pass.TypesInfo.Defs[function.Name].(*types.Func); ok { + context.declarations[object] = function + } + } + } + } + return context +} + +func (context *ps6136ContractContext) named(identity string) *types.Named { + separator := strings.LastIndexByte(identity, '.') + if separator < 1 { + return nil + } + pkg := context.packages[identity[:separator]] + if pkg == nil { + return nil + } + object, ok := pkg.Scope().Lookup(identity[separator+1:]).(*types.TypeName) + if !ok { + return nil + } + named, ok := types.Unalias(object.Type()).(*types.Named) + if !ok || named.TypeParams().Len() != 0 { + return nil + } + return named +} + +func (context *ps6136ContractContext) callable(identity string) *types.Func { + separator := strings.LastIndexByte(identity, '.') + if separator < 1 { + return nil + } + if pkg := context.packages[identity[:separator]]; pkg != nil { + function, _ := pkg.Scope().Lookup(identity[separator+1:]).(*types.Func) + return function + } + named := context.named(identity[:separator]) + if named == nil { + return nil + } + object, _, _ := types.LookupFieldOrMethod(named, true, context.pass.Pkg, identity[separator+1:]) + function, _ := object.(*types.Func) + if function == nil || ps6090FunctionID(function) != identity { + return nil + } + return function +} + +func ps6136FieldVar(typ types.Type, name string) *types.Var { + if pointer, ok := types.Unalias(typ).(*types.Pointer); ok { + typ = pointer.Elem() + } + structure, ok := typ.Underlying().(*types.Struct) + if !ok { + return nil + } + for index := 0; index < structure.NumFields(); index++ { + field := structure.Field(index) + if field.Name() == name { + return field + } + } + return nil +} + +// Build SSA lazily only after a valid contract resolves to this source package. +// Dependency packages provide typed API metadata, never fabricated source bodies. +func ps6136SourcePackage(pass *analysis.Pass) *ssa.Package { + program := ssa.NewProgram(pass.Fset, ssa.SanityCheckFunctions) + seen := make(map[*types.Package]bool) + var add func(*types.Package) + add = func(pkg *types.Package) { + if seen[pkg] { + return + } + seen[pkg] = true + for _, imported := range pkg.Imports() { + add(imported) + } + if pkg != pass.Pkg { + program.CreatePackage(pkg, nil, nil, true) + } + } + add(pass.Pkg) + pkg := program.CreatePackage(pass.Pkg, pass.Files, pass.TypesInfo, true) + pkg.Build() + return pkg +} diff --git a/checks/ps6136_coverage.go b/checks/ps6136_coverage.go new file mode 100644 index 00000000..3c493b66 --- /dev/null +++ b/checks/ps6136_coverage.go @@ -0,0 +1,143 @@ +package checks + +import ( + "go/ast" + "go/token" + "go/types" + + "golang.org/x/tools/go/analysis" +) + +// ps6136ClosedObservations requires a source proof for every workspace field +// observation and every whole-owner exposure in the loaded package. A witness +// in a selected helper cannot establish complete use coverage. Proved nodes +// are supplied only after exact typed argument, extent and ownership checks; +// this routine never interprets configuration as a source proof. +func ps6136ClosedObservations(pass *analysis.Pass, owner *types.Named, workspace *types.Var, proved map[ast.Node]bool) bool { + if owner == nil || workspace == nil { + return false + } + parents := make(map[ast.Node]ast.Node) + for _, file := range pass.Files { + var stack []ast.Node + ast.Inspect(file, func(node ast.Node) bool { + if node == nil { + stack = stack[:len(stack)-1] + return false + } + if len(stack) > 0 { + parents[node] = stack[len(stack)-1] + } + stack = append(stack, node) + return true + }) + } + approved := func(node ast.Node) bool { + for node != nil { + if proved[node] { + return true + } + // A function-wide approval would hide contradictory extra uses. + switch node.(type) { + case *ast.FuncDecl, *ast.FuncLit, *ast.BlockStmt: + return false + } + node = parents[node] + } + return false + } + isOwner := func(typ types.Type) bool { + if pointer, ok := typ.(*types.Pointer); ok { + typ = pointer.Elem() + } + return typ != nil && types.Identical(typ, owner) + } + valid := true + for _, file := range pass.Files { + ast.Inspect(file, func(node ast.Node) bool { + if !valid { + return false + } + switch n := node.(type) { + case *ast.UnaryExpr: + if n.Op == token.AND && !approved(n) { + // An interior struct-field pointer can expose owner storage + // without passing an expression typed *Owner. Slice indexes + // refer to separate backing storage and are not inferred here. + for expression := ps2110Unparen(n.X); expression != nil; { + selector, ok := expression.(*ast.SelectorExpr) + if !ok { + break + } + if isOwner(pass.TypesInfo.TypeOf(selector.X)) { + valid = false + break + } + expression = ps2110Unparen(selector.X) + } + } + case *ast.SelectorExpr: + selection := pass.TypesInfo.Selections[n] + if selection != nil && selection.Obj() == workspace && !approved(n) { + valid = false + } + if isOwner(pass.TypesInfo.TypeOf(n)) && !approved(n) { + outer, ok := parents[n].(*ast.SelectorExpr) + if !ok || outer.X != n { + valid = false + } else if next := pass.TypesInfo.Selections[outer]; next == nil || next.Kind() != types.FieldVal { + valid = false + } + } + case *ast.CompositeLit: + if isOwner(pass.TypesInfo.TypeOf(n)) && !approved(n) { + for _, element := range n.Elts { + keyed, ok := element.(*ast.KeyValueExpr) + if !ok { + valid = false + break + } + identifier, ok := keyed.Key.(*ast.Ident) + if !ok || pass.TypesInfo.Uses[identifier] == workspace && !proved[keyed.Key] { + valid = false + break + } + } + } + case *ast.Ident: + object := pass.TypesInfo.Uses[n] + if _, variable := object.(*types.Var); !variable || !isOwner(object.Type()) || approved(n) { + break + } + parent := parents[n] + for { + parens, ok := parent.(*ast.ParenExpr) + if !ok { + break + } + parent = parents[parens] + } + if comparison, ok := parent.(*ast.BinaryExpr); ok && (comparison.Op == token.EQL || comparison.Op == token.NEQ) { + other := comparison.X + if ps2110Unparen(comparison.X) == n { + other = comparison.Y + } + if identifier, ok := ps2110Unparen(other).(*ast.Ident); ok && pass.TypesInfo.Uses[identifier] == types.Universe.Lookup("nil") { + break // Pointer nil comparison neither aliases nor exposes storage. + } + } + selector, ok := parent.(*ast.SelectorExpr) + if !ok { + valid = false // aliases, captures, sends, opaque arguments, whole stores + break + } + selection := pass.TypesInfo.Selections[selector] + if selection == nil || selection.Kind() != types.FieldVal { + valid = false // implicit-address methods require their exact call proof + } + } + return valid + }) + } + return valid +} diff --git a/checks/ps6136_coverage_test.go b/checks/ps6136_coverage_test.go new file mode 100644 index 00000000..64d80320 --- /dev/null +++ b/checks/ps6136_coverage_test.go @@ -0,0 +1,77 @@ +package checks + +import ( + "go/ast" + "go/parser" + "go/token" + "go/types" + "testing" + + "golang.org/x/tools/go/analysis" +) + +func TestPS6136ClosedObservations(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name, extra string + valid bool + }{ + {"proved_use", "", true}, + {"sibling_owner", `type Other struct{logits []float32};func extra(other *Other){_=other.logits}`, true}, + {"nil_comparison", `func extra(d *Owner)bool{return ((d))==nil || nil!=d}`, true}, + {"pointer_comparison", `func extra(d, other *Owner)bool{return d==other}`, false}, + {"unclassified_use", `func extra(d *Owner){_=d.logits}`, false}, + {"alias", `func extra(d *Owner){alias:=d;_=alias}`, false}, + {"capture", `func extra(d *Owner){fn:=func(){_=d};_=fn}`, false}, + {"interior_geometry_pointer", `func extra(d *Owner){_= &((d.width))}`, false}, + {"opaque_mutation", `func extra(d *Owner){opaque(d)};func opaque(*Owner){}`, false}, + {"whole_reset", `func extra(d *Owner){*d=Owner{}}`, false}, + {"keyed_alternate", `var other=&Owner{logits:[]float32{1}}`, false}, + {"proved_initializer_key", `func initOwner()*Owner{return &Owner{logits:[]float32{1}}}`, true}, + {"proved_initializer_key_not_rhs", `func initOwner(d *Owner)*Owner{return &Owner{logits:d.logits}}`, false}, + {"positional_alternate", `var other=&Owner{[]float32{1},1}`, false}, + {"hidden_method", `type Holder struct{D *Owner};func extra(h Holder){h.D.mutate()};func(d *Owner)mutate(){d.width++}`, false}, + {"hidden_opaque", `type Holder struct{D *Owner};func extra(h Holder){opaque(h.D)};func opaque(*Owner){}`, false}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + fs := token.NewFileSet() + file, err := parser.ParseFile(fs, "coverage.go", `package coverage;type Owner struct{logits []float32;width int};func observed(d *Owner){_=d.logits};`+test.extra, parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + info := &types.Info{Types: map[ast.Expr]types.TypeAndValue{}, Uses: map[*ast.Ident]types.Object{}, Defs: map[*ast.Ident]types.Object{}, Selections: map[*ast.SelectorExpr]*types.Selection{}} + pkg, err := (&types.Config{}).Check("coverage", fs, []*ast.File{file}, info) + if err != nil { + t.Fatal(err) + } + owner := pkg.Scope().Lookup("Owner").Type().(*types.Named) + workspace := owner.Underlying().(*types.Struct).Field(0) + proved := map[ast.Node]bool{} + if test.name == "proved_initializer_key" || test.name == "proved_initializer_key_not_rhs" { + ast.Inspect(file, func(node ast.Node) bool { + if initializer, ok := node.(*ast.KeyValueExpr); ok { + proved[initializer.Key] = true + } + return true + }) + } + for _, declaration := range file.Decls { + function, ok := declaration.(*ast.FuncDecl) + if !ok || function.Name.Name != "observed" { + continue + } + ast.Inspect(function.Body, func(node ast.Node) bool { + if selection, ok := node.(*ast.SelectorExpr); ok { + proved[selection] = true + } + return true + }) + } + pass := &analysis.Pass{Fset: fs, Files: []*ast.File{file}, TypesInfo: info, Pkg: pkg} + if got := ps6136ClosedObservations(pass, owner, workspace, proved); got != test.valid { + t.Fatalf("closed observations=%v want %v", got, test.valid) + } + }) + } +} diff --git a/checks/ps6136_error_barrier.go b/checks/ps6136_error_barrier.go new file mode 100644 index 00000000..533e78bc --- /dev/null +++ b/checks/ps6136_error_barrier.go @@ -0,0 +1,158 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136ConstructorErrorBarrier proves the canonical final publication edge: +// the selected constructor returns this owner only with nil aggregate error; +// the opposite edge calls its exact Release before returning nil/error. This +// does not infer Release effects, prior error-cell invariance, allocator failure +// ownership, or the completeness of earlier construction returns. +func ps6136ConstructorErrorBarrier(context *ps6125SSAContext, owner ps6125SSAReference, ownerType *types.Named, cell ps6125SSAReference, release string) bool { + if context == nil || owner.value == nil || cell.value == nil || ownerType == nil { + return false + } + var published *ssa.Return + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + ret, ok := instruction.(*ssa.Return) + if !ok || len(ret.Results) != 2 { + continue + } + if literal, ok := ret.Results[0].(*ssa.Const); ok && literal.IsNil() { + continue + } + if published != nil || ps6136OwnerRoot(context, ret.Results[0], ownerType) != owner { + return false + } + literal, ok := ret.Results[1].(*ssa.Const) + if !ok || !literal.IsNil() { + return false + } + published = ret + } + } + if published == nil || len(published.Block().Preds) != 1 { + return false + } + // No allocation, write or unknown effect may occur after checking the + // aggregate error and before publication. Only the exact returned owner + // and same error cell's effect-free pointer loads are permitted. + for _, instruction := range published.Block().Instrs { + switch use := instruction.(type) { + case *ssa.DebugRef: + case *ssa.UnOp: + if use.Op != token.MUL || ps6136OwnerRoot(context, use, ownerType) != owner && !ps6136ErrorCellRead(context, use, cell) { + return false + } + case *ssa.Return: + if use != published { + return false + } + default: + return false + } + } + predecessor := published.Block().Preds[0] + if len(predecessor.Instrs) == 0 || len(predecessor.Succs) != 2 { + return false + } + branch, ok := predecessor.Instrs[len(predecessor.Instrs)-1].(*ssa.If) + if !ok { + return false + } + comparison, ok := branch.Cond.(*ssa.BinOp) + if !ok || comparison.Op != token.EQL && comparison.Op != token.NEQ { + return false + } + value, nilValue := comparison.X, comparison.Y + if literal, ok := value.(*ssa.Const); ok && literal.IsNil() { + value, nilValue = nilValue, value + } + literal, ok := nilValue.(*ssa.Const) + if !ok || !literal.IsNil() || !ps6136ErrorCellRead(context, value, cell) { + return false + } + checked := context.reference(value) + load, ok := checked.value.(*ssa.UnOp) + if !ok || checked.context != context || load.Block() != predecessor || comparison.Block() != predecessor { + return false + } + // The predicate must use the current error, not a snapshot taken before + // another allocation or error write. Close the entire read-to-branch + // interval, including precomputed boolean predicates. + reading := false + for _, instruction := range predecessor.Instrs { + if instruction == load { + reading = true + continue + } + if !reading { + continue + } + switch instruction { + case comparison, branch: + default: + if _, debug := instruction.(*ssa.DebugRef); !debug { + return false + } + } + } + if !reading { + return false + } + errorIndex := 1 + if comparison.Op == token.NEQ { + errorIndex = 0 + } + if predecessor.Succs[1-errorIndex] != published.Block() { + return false + } + failure := predecessor.Succs[errorIndex] + var cleanup *ssa.Call + var returned *ssa.Return + for _, instruction := range failure.Instrs { + switch use := instruction.(type) { + case *ssa.Call: + function := use.Call.StaticCallee() + if function == nil || function.Object() == nil || use.Call.IsInvoke() || len(use.Call.Args) != 1 || cleanup != nil { + return false + } + object, ok := function.Object().(*types.Func) + if !ok || ps6090FunctionID(object) != release || ps6136OwnerRoot(context, use.Call.Args[0], ownerType) != owner { + return false + } + cleanup = use + case *ssa.Return: + if returned != nil || len(use.Results) != 2 { + return false + } + literal, ok := use.Results[0].(*ssa.Const) + if !ok || !literal.IsNil() || !ps6136ErrorCellRead(context, use.Results[1], cell) { + return false + } + returned = use + case *ssa.DebugRef: + case *ssa.UnOp: + if use.Op != token.MUL || ps6136OwnerRoot(context, use, ownerType) != owner && !ps6136ErrorCellRead(context, use, cell) { + return false + } + default: + return false + } + } + return cleanup != nil && returned != nil && context.flow.instructionDominates(cleanup, returned) +} + +func ps6136ErrorCellRead(context *ps6125SSAContext, value ssa.Value, cell ps6125SSAReference) bool { + reference := context.reference(value) + load, ok := reference.value.(*ssa.UnOp) + return ok && load.Op == token.MUL && reference.context.reference(load.X) == cell +} diff --git a/checks/ps6136_error_barrier_test.go b/checks/ps6136_error_barrier_test.go new file mode 100644 index 00000000..1f215fd4 --- /dev/null +++ b/checks/ps6136_error_barrier_test.go @@ -0,0 +1,73 @@ +package checks + +import ( + "go/types" + "strings" + "testing" + + "golang.org/x/tools/go/ssa" +) + +func TestPS6136ConstructorFailurePublicationEdges(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name, barrier string + want bool + }{ + {"canonical", "if err!=nil{d.Release();return nil,err};return d,nil", true}, + {"equivalent", "if nil!=err{d.Release();return nil,err};return d,nil", true}, + {"no-release", "if err!=nil{return nil,err};return d,nil", false}, + {"wrong-release", "if err!=nil{d.Other();return nil,err};return d,nil", false}, + {"wrong-owner", "if err!=nil{(&owner{}).Release();return nil,err};return d,nil", false}, + {"publish-failure", "if err!=nil{d.Release();return d,err};return d,nil", false}, + {"discard-error", "if err!=nil{d.Release();return nil,nil};return d,nil", false}, + {"unguarded", "return d,nil", false}, + {"late-allocation", "if err!=nil{d.Release();return nil,err};mk();return d,nil", false}, + {"late-error-write", "if err!=nil{d.Release();return nil,err};err=native();return d,nil", false}, + {"late-error-reset", "if err!=nil{d.Release();return nil,err};err=nil;return d,nil", false}, + {"stale-error-load", "old:=err;mk();if old!=nil{d.Release();return nil,err};return d,nil", false}, + {"stale-error-condition", "bad:=err!=nil;mk();if bad{d.Release();return nil,err};return d,nil", false}, + {"wrong-state", "var other error;if other!=nil{d.Release();return nil,other};return d,nil", false}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + pkg := ps6125TestSSA(t, strings.Replace(`package extent +type owner struct{} +func(*owner)Release(){} +func(*owner)Other(){} +func native()error{return nil} +func construct()(*owner,error){ + d:=&owner{};var err error + mk:=func(){e:=native();if e!=nil{err=e}} + mk() + _ = err + BARRIER +}`, "BARRIER", test.barrier, 1)) + function := pkg.Func("construct") + context := ps6125NewSSAContext(function, nil, nil, 64) + ownerType := pkg.Pkg.Scope().Lookup("owner").Type().(*types.Named) + var owner, cell ps6125SSAReference + for _, block := range function.Blocks { + for _, instruction := range block.Instrs { + allocation, ok := instruction.(*ssa.Alloc) + if !ok { + continue + } + if types.Identical(allocation.Type(), types.NewPointer(ownerType)) && owner.value == nil { + owner = context.reference(allocation) + } + if types.Identical(allocation.Type(), types.NewPointer(types.Universe.Lookup("error").Type())) && allocation.Comment == "err" { + cell = context.reference(allocation) + } + } + } + method, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, "Release") + if owner.value == nil || cell.value == nil { + t.Fatal("fixture source identities unavailable") + } + if got := ps6136ConstructorErrorBarrier(context, owner, ownerType, cell, ps6090FunctionID(method.(*types.Func))); got != test.want { + t.Fatalf("publication edge %v, want %v", got, test.want) + } + }) + } +} diff --git a/checks/ps6136_extents.go b/checks/ps6136_extents.go new file mode 100644 index 00000000..489fd09b --- /dev/null +++ b/checks/ps6136_extents.go @@ -0,0 +1,103 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136Extents adds only source-proved immutable owner geometry to the +// invocation-specialized PS6125 algebra. It does not infer field invariance +// from repeated loads; callers must close owner mutation/exposure first. +type ps6136Extents struct { + owner ps6125SSAReference + ownerType *types.Named + immutable map[*types.Var]ps6125Extent + cache map[ps6125SSAReference]ps6125Extent + active map[ps6125SSAReference]bool + remaining int +} + +func (facts *ps6136Extents) extent(context *ps6125SSAContext, value ssa.Value) ps6125Extent { + if context == nil || value == nil || facts.remaining <= 0 || facts.owner.value == nil { + return ps6125Extent{} + } + if scalar := context.scalar(value); scalar.state == ps6125Integer { + return scalar.extent + } + reference := context.reference(value) + if reference.value == nil { + return ps6125Extent{} + } + if reference.context != context || reference.value != value { + return facts.extent(reference.context, reference.value) + } + if facts.cache == nil { + facts.cache = make(map[ps6125SSAReference]ps6125Extent) + facts.active = make(map[ps6125SSAReference]bool) + } + if result, found := facts.cache[reference]; found { + return result + } + if facts.active[reference] { + return ps6125Extent{} + } + facts.remaining-- + facts.active[reference] = true + var result ps6125Extent + if types.Identical(value.Type(), types.Typ[types.Int]) { + paths := ps6125AccessPaths{flow: context.flow} + path := paths.resolve(value) + root := context.reference(path.access.root) + if facts.ownerType != nil { + root = ps6136AccessOwnerRoot(context, path, facts.ownerType) + } + if path.known && len(path.access.fields) == 1 && root == facts.owner { + result = facts.immutable[path.access.fields[0]] + } + } + if !result.known { + switch instruction := value.(type) { + case *ssa.BinOp: + if instruction.Op == token.MUL && types.Identical(instruction.Type(), types.Typ[types.Int]) { + result = ps6125MultiplyExtents(facts.extent(context, instruction.X), facts.extent(context, instruction.Y)) + } + case *ssa.Phi: + found := false + for index, predecessor := range instruction.Block().Preds { + if !context.flow.edges[ps6125SSAEdge{from: predecessor, to: instruction.Block()}] { + continue + } + incoming := facts.extent(context, instruction.Edges[index]) + if !incoming.known || found && !ps6125SameExtent(result, incoming) { + result = ps6125Extent{} + break + } + result, found = incoming, true + } + case *ssa.Call: + if builtin, ok := instruction.Call.Value.(*ssa.Builtin); ok && builtin.Name() == "len" && len(instruction.Call.Args) == 1 { + result = facts.length(context, instruction.Call.Args[0]) + } + } + } + delete(facts.active, reference) + facts.cache[reference] = result + return result +} + +func (facts *ps6136Extents) length(context *ps6125SSAContext, value ssa.Value) ps6125Extent { + if context == nil || value == nil { + return ps6125Extent{} + } + if length := context.length(value); length.known { + return length + } + reference := context.reference(value) + allocation, ok := reference.value.(*ssa.MakeSlice) + if !ok { + return ps6125Extent{} + } + return facts.extent(reference.context, allocation.Len) +} diff --git a/checks/ps6136_extents_test.go b/checks/ps6136_extents_test.go new file mode 100644 index 00000000..b1788bb6 --- /dev/null +++ b/checks/ps6136_extents_test.go @@ -0,0 +1,67 @@ +package checks + +import ( + "go/types" + "testing" + + "golang.org/x/tools/go/ssa" +) + +// Unit tests for the extent layer, not a substitute for complete owner replay. +func TestPS6136InvocationExtents(t *testing.T) { + t.Parallel() + pkg := ps6125TestSSA(t, `package extent +type owner struct{width int; other int} +func leaf([]float32){} +func work(d,sibling *owner,tokens []int,lastOnly bool){ + rows:=len(tokens);if lastOnly{rows=1} + leaf(make([]float32,rows*d.width)) + leaf(make([]float32,rows*sibling.width)) + leaf(make([]float32,rows*d.other)) +} +func common(d,sibling *owner,tokens []int){work(d,sibling,tokens,true)} +func bulk(d,sibling *owner,tokens []int){work(d,sibling,tokens,false)} +`) + for _, name := range []string{"common", "bulk"} { + t.Run(name, func(t *testing.T) { + t.Parallel() + entry := pkg.Func(name) + var pool ps6125ExtentPool + structure := pkg.Pkg.Scope().Lookup("owner").Type().Underlying().(*types.Struct) + width := ps6125SymbolicExtent(pool.identity(entry.Params[0].Object(), []*types.Var{structure.Field(0)}, false)) + rows := ps6125SymbolicExtent(pool.identity(entry.Params[2].Object(), nil, true)) + root := ps6125NewSSAContext(entry, nil, map[*ssa.Parameter]ps6125Extent{entry.Params[2]: rows}, 8) + child := root.call(ps6125ContextCalls(root)[0]) + if child == nil { + t.Fatal("exact helper invocation unavailable") + } + facts := ps6136Extents{owner: root.reference(entry.Params[0]), immutable: map[*types.Var]ps6125Extent{structure.Field(0): width}, remaining: 1000} + var calls []*ssa.Call + for _, call := range ps6125ContextCalls(child) { + if call.Call.StaticCallee() == pkg.Func("leaf") { + calls = append(calls, call) + } + } + if len(calls) != 3 { + t.Fatalf("leaf count %d", len(calls)) + } + active := rows + if name == "common" { + active = ps6125ConstantExtent(1) + } + want := ps6125MultiplyExtents(active, width) + if got := facts.length(child, calls[0].Call.Args[0]); !ps6125SameExtent(got, want) { + t.Fatal("actual owner specialized host extent unavailable") + } + for _, call := range calls[1:] { + if facts.length(child, call.Call.Args[0]).known { + t.Fatal("sibling owner/unknown field became immutable geometry") + } + } + facts.remaining = 0 + if facts.length(child, calls[0].Call.Args[0]).known { + t.Fatal("exhausted proof budget accepted") + } + }) + } +} diff --git a/checks/ps6136_external_owner_test.go b/checks/ps6136_external_owner_test.go new file mode 100644 index 00000000..636b4dbc --- /dev/null +++ b/checks/ps6136_external_owner_test.go @@ -0,0 +1,46 @@ +package checks + +import ( + "bytes" + "go/types" + "os" + "path/filepath" + "testing" + + "golang.org/x/tools/go/packages" +) + +// Optional read-only replay against a complete pinned public checkout. The +// ordinary hermetic fixtures remain mandatory; this additional audit neither +// downloads source nor executes a GPU backend during the repository's gates. +func TestPS6136PinnedCompletePackageConstructorInventory(t *testing.T) { + directory := os.Getenv("PERFSCAN_PS6136_OWNER_DIR") + if directory == "" { + t.Skip("set PERFSCAN_PS6136_OWNER_DIR to the exact before-owner checkout") + } + for _, name := range []string{"gpt.go", "decoder.go"} { + actual, err := os.ReadFile(filepath.Join(directory, "llamagpu", name)) + if err != nil { + t.Fatal(err) + } + pinned, err := os.ReadFile("testdata/ps6136-owner/before/" + name + ".txt") + if err != nil || !bytes.Equal(actual, pinned) { + t.Fatalf("%s is not the pinned before source", name) + } + } + loaded, err := packages.Load(&packages.Config{Dir: directory, Mode: packages.NeedName | packages.NeedFiles | packages.NeedCompiledGoFiles | packages.NeedSyntax | packages.NeedTypes | packages.NeedTypesInfo | packages.NeedImports | packages.NeedDeps}, "./llamagpu") + if err != nil || packages.PrintErrors(loaded) != 0 || len(loaded) != 1 { + t.Fatal("complete owner package does not load cleanly") + } + source := loaded[0] + fixture := &analysisOwnerFixture{source.Fset, source.Syntax, source.TypesInfo, source.Types, nil} + pkg := ps6136FixtureSSA(fixture) + for _, name := range []string{"GPTDecoder", "Decoder"} { + owner := source.Types.Scope().Lookup(name).Type().(*types.Named) + inventory := ps6136ConstructorInventory(pkg, owner) + if inventory == nil { + t.Fatalf("complete real package %s constructor inventory unknown", name) + } + t.Logf("complete pinned native package: %s %d fresh constructor functions", name, len(inventory)) + } +} diff --git a/checks/ps6136_factory.go b/checks/ps6136_factory.go new file mode 100644 index 00000000..c02d67a4 --- /dev/null +++ b/checks/ps6136_factory.go @@ -0,0 +1,80 @@ +package checks + +import ( + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136FactoryResult verifies actual input and returned-slot provenance, not +// successful allocation or native ownership. The caller must additionally +// prove same-owner retention, typed backend binding, error cleanup and release. +func ps6136FactoryResult(context *ps6125SSAContext, input ps6125SSAReference, allocator, slot *types.Var) *ssa.Call { + if context == nil || input.value == nil || allocator == nil || slot == nil { + return nil + } + signature, ok := allocator.Type().Underlying().(*types.Signature) + if !ok || signature.Variadic() || signature.TypeParams().Len() != 0 || signature.Params().Len() != 1 || signature.Results().Len() != 2 || + !types.Identical(signature.Params().At(0).Type(), types.NewSlice(types.Typ[types.Float32])) || + !types.Identical(signature.Results().At(0).Type(), slot.Type()) || + !types.Identical(signature.Results().At(1).Type(), types.Universe.Lookup("error").Type()) { + return nil + } + paths := ps6125AccessPaths{flow: context.flow} + var backend *ssa.Call + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + call, ok := instruction.(*ssa.Call) + if !ok { + continue + } + path := paths.resolve(call.Call.Value) + if !path.known || len(path.access.fields) == 0 || path.access.fields[len(path.access.fields)-1] != allocator { + continue + } + if backend != nil || call.Call.IsInvoke() || len(call.Call.Args) != 1 || context.reference(call.Call.Args[0]) != input { + return nil + } + backend = call + } + } + if backend == nil { + return nil + } + populated, empty := 0, 0 + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + returned, ok := instruction.(*ssa.Return) + if !ok { + continue + } + if len(returned.Results) != 1 { + return nil + } + _, fields, known := context.returnedFields(returned, 0) + if !known { + return nil + } + reference := fields[slot] + if reference.value == nil { + empty++ + continue + } + extracted, ok := reference.value.(*ssa.Extract) + if !ok || reference.context != context || extracted.Tuple != backend || extracted.Index != 0 { + return nil + } + populated++ + } + } + if populated != 1 || empty != 2 { + return nil + } + return backend +} diff --git a/checks/ps6136_factory_error.go b/checks/ps6136_factory_error.go new file mode 100644 index 00000000..2363531d --- /dev/null +++ b/checks/ps6136_factory_error.go @@ -0,0 +1,94 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136FactoryErrorCell proves that the exact allocator error is propagated +// into one captured constructor error cell before its failed slot return. It +// does not establish the constructor's publication/cleanup barrier or native +// failure ownership. Those are separate mandatory lifetime gates. +func ps6136FactoryErrorCell(context *ps6125SSAContext, backend *ssa.Call) ps6125SSAReference { + if context == nil || backend == nil || backend.Parent() != context.flow.function { + return ps6125SSAReference{} + } + var write *ssa.Store + var failed *ssa.Return + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + if store, ok := instruction.(*ssa.Store); ok { + extracted, ok := store.Val.(*ssa.Extract) + if !ok || extracted.Tuple != backend || extracted.Index != 1 { + continue + } + if write != nil || !types.Identical(store.Addr.Type(), types.NewPointer(types.Universe.Lookup("error").Type())) { + return ps6125SSAReference{} + } + write = store + } + if ret, ok := instruction.(*ssa.Return); ok && ps6136NativeErrorBranch(context, ret, backend, true) { + if failed != nil { + return ps6125SSAReference{} + } + failed = ret + } + } + } + if write == nil || failed == nil || !context.flow.instructionDominates(write, failed) { + return ps6125SSAReference{} + } + cell := ps6136ErrorCellRoot(context, write.Addr) + allocation, ok := cell.value.(*ssa.Alloc) + if !ok || allocation.Parent() != cell.context.flow.function || cell.context == context { + return ps6125SSAReference{} + } + return cell +} + +// Resolve only immutable pointer-to-error captures, not mutable error values. +func ps6136ErrorCellRoot(context *ps6125SSAContext, value ssa.Value) ps6125SSAReference { + if context == nil || value == nil || !types.Identical(value.Type(), types.NewPointer(types.Universe.Lookup("error").Type())) { + return ps6125SSAReference{} + } + reference := context.reference(value) + seen := make(map[ps6125SSAReference]bool) + for reference.value != nil && !seen[reference] { + seen[reference] = true + if _, ok := reference.value.(*ssa.Alloc); ok { + return reference + } + load, ok := reference.value.(*ssa.UnOp) + if !ok || load.Op != token.MUL { + break + } + cell := reference.context.reference(load.X) + allocation, ok := cell.value.(*ssa.Alloc) + if !ok || !types.Identical(allocation.Type(), types.NewPointer(value.Type())) { + break + } + store := ps6136ClosedOwnerCell(cell.context, allocation) + if store == nil { + break + } + var point ssa.Instruction = load + current := reference.context + for current != cell.context && current.parent != nil { + point = current.site + current = current.parent + } + if current != cell.context { + point = ps6136ReturnedCapturePoint(reference.context, cell) + } + if point == nil || !cell.context.flow.instructionDominates(store, point) { + break + } + reference = cell.context.reference(store.Val) + } + return ps6125SSAReference{} +} diff --git a/checks/ps6136_field_effects.go b/checks/ps6136_field_effects.go new file mode 100644 index 00000000..02ab8a6f --- /dev/null +++ b/checks/ps6136_field_effects.go @@ -0,0 +1,104 @@ +package checks + +import ( + "go/ast" + "go/token" + "go/types" + + "golang.org/x/tools/go/analysis" +) + +// ps6136FieldEffects closes direct writes/addresses of one exact typed field. +// Only separately proved fresh-constructor initialization targets may be +// allowed. This is not owner-alias closure: ps6136ClosedObservations must also +// reject indirect receiver/holder/opaque mutation paths. A write to a sibling +// child member is not a whole-field write (ProfileMetalStep's recorder changes +// must not falsely imply its allocator field changed). +func ps6136FieldEffects(pass *analysis.Pass, field *types.Var, initializers map[ast.Node]bool) bool { + if pass == nil || field == nil { + return false + } + parents := make(map[ast.Node]ast.Node) + var occurrences []ast.Node + for _, file := range pass.Files { + var stack []ast.Node + ast.Inspect(file, func(node ast.Node) bool { + if node == nil { + stack = stack[:len(stack)-1] + return false + } + if len(stack) > 0 { + parents[node] = stack[len(stack)-1] + } + stack = append(stack, node) + if selector, ok := node.(*ast.SelectorExpr); ok { + if selection := pass.TypesInfo.Selections[selector]; selection != nil && selection.Obj() == field { + occurrences = append(occurrences, selector) + } + } + if pair, ok := node.(*ast.KeyValueExpr); ok { + if key, ok := pair.Key.(*ast.Ident); ok && pass.TypesInfo.Uses[key] == field { + occurrences = append(occurrences, pair) + } + } + return true + }) + } + for _, occurrence := range occurrences { + if initializers[occurrence] { + continue + } + if _, literal := occurrence.(*ast.KeyValueExpr); literal { + return false + } + target := occurrence + // An interior address of a child member also exposes the containing + // field's storage; sibling assignments alone do not. + for child := occurrence; child != nil; { + parent := parents[child] + if address, ok := parent.(*ast.UnaryExpr); ok && address.Op == token.AND && address.X == child { + return false + } + switch parent := parent.(type) { + case *ast.ParenExpr: + child = parent + case *ast.SelectorExpr: + if parent.X != child { + child = nil + } else { + child = parent + } + default: + child = nil + } + } + for { + parens, ok := parents[target].(*ast.ParenExpr) + if !ok { + break + } + target = parens + } + switch parent := parents[target].(type) { + case *ast.UnaryExpr: + if parent.Op == token.AND && parent.X == target { + return false + } + case *ast.AssignStmt: + for _, left := range parent.Lhs { + if left == target { + return false + } + } + case *ast.IncDecStmt: + if parent.X == target { + return false + } + case *ast.RangeStmt: + if parent.Tok == token.ASSIGN && (parent.Key == target || parent.Value == target) { + return false + } + } + } + return true +} diff --git a/checks/ps6136_field_effects_test.go b/checks/ps6136_field_effects_test.go new file mode 100644 index 00000000..4c498c2c --- /dev/null +++ b/checks/ps6136_field_effects_test.go @@ -0,0 +1,54 @@ +package checks + +import ( + "go/ast" + "go/parser" + "go/token" + "go/types" + "testing" + + "golang.org/x/tools/go/analysis" +) + +func TestPS6136TypedFieldEffects(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name, body, field string + want bool + }{ + {"read", "_=d.width", "width", true}, + {"write", "d.width=0", "width", false}, + {"increment", "d.width++", "width", false}, + {"range-rebind", "for d.width=range 3{}", "width", false}, + {"parenthesized-address", "_= &((d.width))", "width", false}, + {"same-spelling-other-type", "o:=Other{};o.width=0;_=d.width", "width", true}, + {"whole-ops", "d.ops=Ops{}", "ops", false}, + {"sibling-recorder", "d.ops.recorder=nil", "ops", true}, + {"interior-ops-address", "_= &((d.ops).recorder)", "ops", false}, + {"allocator-write", "d.ops.allocator=nil", "allocator", false}, + {"recorder-not-allocator", "d.ops.recorder=nil", "allocator", true}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + fs := token.NewFileSet() + file, err := parser.ParseFile(fs, "effects.go", `package effects;type Ops struct{allocator,recorder func()};type Owner struct{width int;ops Ops};type Other struct{width int};func f(d *Owner){`+test.body+`}`, parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + info := &types.Info{Types: map[ast.Expr]types.TypeAndValue{}, Uses: map[*ast.Ident]types.Object{}, Defs: map[*ast.Ident]types.Object{}, Selections: map[*ast.SelectorExpr]*types.Selection{}} + pkg, err := (&types.Config{}).Check("effects", fs, []*ast.File{file}, info) + if err != nil { + t.Fatal(err) + } + owner := pkg.Scope().Lookup("Owner").Type() + if test.field == "allocator" { + owner = pkg.Scope().Lookup("Ops").Type() + } + field, _, _ := types.LookupFieldOrMethod(owner, true, pkg, test.field) + pass := &analysis.Pass{Fset: fs, Files: []*ast.File{file}, Pkg: pkg, TypesInfo: info} + if got := ps6136FieldEffects(pass, field.(*types.Var), nil); got != test.want { + t.Fatalf("field effects %v, want %v", got, test.want) + } + }) + } +} diff --git a/checks/ps6136_forwarding.go b/checks/ps6136_forwarding.go new file mode 100644 index 00000000..273e4bd8 --- /dev/null +++ b/checks/ps6136_forwarding.go @@ -0,0 +1,172 @@ +package checks + +import ( + "go/ast" + "go/types" + "slices" + + "golang.org/x/tools/go/analysis" +) + +// The selected F32 projector lowers its exact input/output, active rows and +// source-established stored weight dimensions to the configured recorder API. +func ps6136ProjectionLowering(pass *analysis.Pass, declaration *ast.FuncDecl, methodID string, weight, inner, width *types.Var) bool { + if pass == nil || declaration == nil || declaration.Body == nil || declaration.Recv == nil || len(declaration.Recv.List) != 1 || len(declaration.Recv.List[0].Names) != 1 || len(declaration.Body.List) != 1 { + return false + } + function, ok := pass.TypesInfo.Defs[declaration.Name].(*types.Func) + if !ok { + return false + } + signature, ok := function.Type().(*types.Signature) + if !ok || signature.Variadic() || signature.TypeParams().Len() != 0 || signature.Params().Len() != 4 || signature.Results().Len() != 1 || !types.Identical(signature.Results().At(0).Type(), types.Universe.Lookup("error").Type()) || !types.Identical(signature.Params().At(3).Type(), types.Typ[types.Int]) { + return false + } + returned, ok := declaration.Body.List[0].(*ast.ReturnStmt) + if !ok || len(returned.Results) != 1 { + return false + } + call, ok := returned.Results[0].(*ast.CallExpr) + if !ok || call.Ellipsis.IsValid() || len(call.Args) != 6 { + return false + } + callee, target, ok := typedCallee(pass, call.Fun) + if !ok || ps6090FunctionID(callee) != methodID || target.Variadic() || target.TypeParams().Len() != 0 || target.Params().Len() != 6 || target.Results().Len() != 1 || !types.Identical(target.Results().At(0).Type(), types.Universe.Lookup("error").Type()) { + return false + } + selector, ok := call.Fun.(*ast.SelectorExpr) + if !ok { + return false + } + identifier, ok := selector.X.(*ast.Ident) + if !ok || pass.TypesInfo.Uses[identifier] != signature.Params().At(0) { + return false + } + parameters := [6]int{1, -1, 2, 3, -1, -1} + fields := [6]*types.Var{nil, weight, nil, nil, inner, width} + receiver := pass.TypesInfo.Defs[declaration.Recv.List[0].Names[0]] + for index, argument := range call.Args { + if !types.Identical(pass.TypesInfo.TypeOf(argument), target.Params().At(index).Type()) { + return false + } + if parameter := parameters[index]; parameter >= 0 { + identifier, ok := argument.(*ast.Ident) + if !ok || pass.TypesInfo.Uses[identifier] != signature.Params().At(parameter) { + return false + } + continue + } + selector, ok := argument.(*ast.SelectorExpr) + if !ok { + return false + } + selection := pass.TypesInfo.Selections[selector] + identifier, ok := selector.X.(*ast.Ident) + if !ok || pass.TypesInfo.Uses[identifier] != receiver || selection == nil || selection.Kind() != types.FieldVal || selection.Obj() != fields[index] { + return false + } + } + return true +} + +// ps6136ForwardedLeaf proves a direct typed adapter call preserves every +// buffer/geometry argument in order. Native dimension/access semantics and the +// native-recorder factory binding are separate facts, never inferred from a +// method name. The buffer bridge's source identity/ownership must be closed by +// ps6136BufferBridge and constructor result binding, respectively. +func ps6136ForwardedLeaf(pass *analysis.Pass, declaration *ast.FuncDecl, nativeID string, bridge *types.Func, buffers []int) bool { + if pass == nil || declaration == nil || declaration.Body == nil || declaration.Recv == nil || len(declaration.Recv.List) != 1 || len(declaration.Recv.List[0].Names) != 1 || len(declaration.Body.List) != 1 || bridge == nil { + return false + } + function, ok := pass.TypesInfo.Defs[declaration.Name].(*types.Func) + if !ok { + return false + } + signature, ok := function.Type().(*types.Signature) + if !ok || signature.Variadic() || signature.TypeParams().Len() != 0 || signature.Results().Len() != 1 || !types.Identical(signature.Results().At(0).Type(), types.Universe.Lookup("error").Type()) { + return false + } + for _, index := range buffers { + if index < 0 || index >= signature.Params().Len() { + return false + } + } + returned, ok := declaration.Body.List[0].(*ast.ReturnStmt) + if !ok || len(returned.Results) != 1 { + return false + } + call, ok := returned.Results[0].(*ast.CallExpr) + if !ok || call.Ellipsis.IsValid() || len(call.Args) != signature.Params().Len() { + return false + } + native, nativeSignature, ok := typedCallee(pass, call.Fun) + if !ok || ps6090FunctionID(native) != nativeID || nativeSignature.Recv() == nil || nativeSignature.Variadic() || nativeSignature.TypeParams().Len() != 0 || nativeSignature.Params().Len() != len(call.Args) || nativeSignature.Results().Len() != 1 || !types.Identical(nativeSignature.Results().At(0).Type(), types.Universe.Lookup("error").Type()) { + return false + } + method, ok := call.Fun.(*ast.SelectorExpr) + if !ok { + return false + } + field, ok := method.X.(*ast.SelectorExpr) + if !ok { + return false + } + receiver, ok := field.X.(*ast.Ident) + if !ok || pass.TypesInfo.Uses[receiver] != pass.TypesInfo.Defs[declaration.Recv.List[0].Names[0]] { + return false + } + selection := pass.TypesInfo.Selections[field] + if selection == nil || selection.Kind() != types.FieldVal || !types.Identical(selection.Type(), nativeSignature.Recv().Type()) { + return false + } + for index, argument := range call.Args { + if slices.Contains(buffers, index) { + conversion, ok := argument.(*ast.CallExpr) + if !ok || conversion.Ellipsis.IsValid() || len(conversion.Args) != 1 { + return false + } + converter, conversionSignature, ok := typedCallee(pass, conversion.Fun) + if !ok || converter != bridge || conversionSignature.Variadic() || conversionSignature.Params().Len() != 1 || conversionSignature.Results().Len() != 1 || !types.Identical(conversionSignature.Params().At(0).Type(), signature.Params().At(index).Type()) || !types.Identical(conversionSignature.Results().At(0).Type(), nativeSignature.Params().At(index).Type()) { + return false + } + argument = conversion.Args[0] + } else if !types.Identical(signature.Params().At(index).Type(), types.Typ[types.Int]) || !types.Identical(nativeSignature.Params().At(index).Type(), types.Typ[types.Int]) { + return false + } + identifier, ok := argument.(*ast.Ident) + if !ok || pass.TypesInfo.Uses[identifier] != signature.Params().At(index) { + return false + } + } + return true +} + +// ps6136BufferBridge proves the source assertion unwraps the exact constructor +// wrapper and native buffer field, without subviews, pointer offsets or helpers. +func ps6136BufferBridge(pass *analysis.Pass, declaration *ast.FuncDecl, wrapper *types.Named, field *types.Var) bool { + if pass == nil || declaration == nil || declaration.Body == nil || len(declaration.Body.List) != 1 || wrapper == nil || field == nil { + return false + } + function, ok := pass.TypesInfo.Defs[declaration.Name].(*types.Func) + if !ok { + return false + } + signature, ok := function.Type().(*types.Signature) + if !ok || signature.Recv() != nil || signature.Variadic() || signature.TypeParams().Len() != 0 || signature.Params().Len() != 1 || signature.Results().Len() != 1 || !types.Identical(signature.Results().At(0).Type(), field.Type()) { + return false + } + returned, ok := declaration.Body.List[0].(*ast.ReturnStmt) + if !ok || len(returned.Results) != 1 { + return false + } + selector, ok := returned.Results[0].(*ast.SelectorExpr) + if !ok || pass.TypesInfo.Selections[selector] == nil || pass.TypesInfo.Selections[selector].Obj() != field { + return false + } + assertion, ok := selector.X.(*ast.TypeAssertExpr) + if !ok || assertion.Type == nil || !types.Identical(pass.TypesInfo.TypeOf(assertion.Type), wrapper) { + return false + } + identifier, ok := assertion.X.(*ast.Ident) + return ok && pass.TypesInfo.Uses[identifier] == signature.Params().At(0) +} diff --git a/checks/ps6136_forwarding_test.go b/checks/ps6136_forwarding_test.go new file mode 100644 index 00000000..1d04f99d --- /dev/null +++ b/checks/ps6136_forwarding_test.go @@ -0,0 +1,61 @@ +package checks + +import ( + "go/ast" + "go/types" + "testing" + + "golang.org/x/tools/go/analysis" +) + +func TestPS6136AuthenticNativeLeafForwarding(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, "before", true, true) + if err != nil { + t.Fatal(err) + } + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg} + bridge := fixture.pkg.Scope().Lookup("mb").(*types.Func) + wrapper := fixture.pkg.Scope().Lookup("mBuf").Type().(*types.Named) + field := wrapper.Underlying().(*types.Struct).Field(0) + count := 0 + for _, file := range fixture.files { + for _, declaration := range file.Decls { + function, ok := declaration.(*ast.FuncDecl) + if !ok { + continue + } + if function.Name.Name == "mb" { + if !ps6136BufferBridge(pass, function, wrapper, field) { + t.Fatal("actual Metal buffer adapter source assertion not proved") + } + count++ + } + if object, ok := fixture.info.Defs[function.Name].(*types.Func); ok && ps6090FunctionID(object) == "github.com/jxsl13/goai/llamagpu.f32Linear.record" { + projector := fixture.pkg.Scope().Lookup("f32Linear").Type() + field := func(name string) *types.Var { + object, _, _ := types.LookupFieldOrMethod(projector, true, fixture.pkg, name) + return object.(*types.Var) + } + if !ps6136ProjectionLowering(pass, function, "github.com/jxsl13/goai/llamagpu.recorder.MatMul", field("w"), field("k"), field("n")) { + t.Fatal("actual F32 projection buffer/rows/width lowering not proved") + } + count++ + } + if function.Recv == nil || function.Name.Name != "MatMul" && function.Name.Name != "AddBias" { + continue + } + id := "github.com/jxsl13/goai/backend/metal.Recorder." + function.Name.Name + if !ps6136ForwardedLeaf(pass, function, id, bridge, []int{0, 1, 2}) { + t.Fatal("actual Native Metal buffer + geometry role-preserving adapter not proved") + } + if ps6136ForwardedLeaf(pass, function, "other.Recorder."+function.Name.Name, bridge, []int{0, 1, 2}) { + t.Fatal("wrong native identity accepted") + } + count++ + } + } + if count != 4 { + t.Fatalf("actual native source bridges %d", count) + } +} diff --git a/checks/ps6136_head_owner_test.go b/checks/ps6136_head_owner_test.go new file mode 100644 index 00000000..c791209a --- /dev/null +++ b/checks/ps6136_head_owner_test.go @@ -0,0 +1,77 @@ +package checks + +import ( + "go/types" + "testing" + + "golang.org/x/tools/go/ssa" +) + +func TestPS6136AuthenticModelHeadWidth(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + for _, name := range []string{"newGPTDecoder", "newDecoder"} { + t.Run(revision+"/"+name, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, revision, true, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + function := pkg.Func(name) + entryName := "NewGPT" + if name == "newDecoder" { + entryName = "New" + } + entry := ps6125NewSSAContext(pkg.Func(entryName), nil, nil, 512) + var context *ps6125SSAContext + for _, call := range ps6125ContextCalls(entry) { + if call.Call.StaticCallee() == function { + context = entry.call(call) + } + } + if context == nil { + t.Fatal("selected public constructor source invocation unavailable") + } + modelType := function.Params[0].Type().Underlying().(*types.Pointer).Elem().(*types.Named) + headName, projectorName := "Head", "head" + ownerName := "GPTDecoder" + if name == "newDecoder" { + headName, projectorName, ownerName = "Out", "out", "Decoder" + } + head, _, _ := types.LookupFieldOrMethod(modelType, true, modelType.Obj().Pkg(), headName) + tensorType := head.Type().Underlying().(*types.Pointer).Elem().(*types.Named) + shape, _, _ := types.LookupFieldOrMethod(tensorType, true, tensorType.Obj().Pkg(), "Shape") + ownerType := pkg.Pkg.Scope().Lookup(ownerName).Type().(*types.Named) + projector, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, projectorName) + projectors := make(map[*types.Named]*types.Var) + for _, kind := range []string{"f32Linear"} { + named := pkg.Pkg.Scope().Lookup(kind).Type().(*types.Named) + width, _, _ := types.LookupFieldOrMethod(named, true, pkg.Pkg, "n") + projectors[named] = width.(*types.Var) + } + paths := ps6125AccessPaths{flow: context.flow} + count := 0 + for _, block := range function.Blocks { + for _, instruction := range block.Instrs { + store, ok := instruction.(*ssa.Store) + if !ok { + continue + } + path := paths.resolve(store.Addr) + if !path.known || len(path.access.fields) != 1 || path.access.fields[0] != projector { + continue + } + if !ps6136ProjectorHeadWidth(context, store.Val, projectors, context.reference(function.Params[0]), modelType, head.(*types.Var), ps6090FunctionID(shape.(*types.Func)), 1, 256) { + t.Fatal("actual selected model head -> source linear factory -> projector output width not proved") + } + count++ + } + } + if count != 1 { + t.Fatalf("projector assignments %d", count) + } + }) + } + } +} diff --git a/checks/ps6136_head_width.go b/checks/ps6136_head_width.go new file mode 100644 index 00000000..26d3ad86 --- /dev/null +++ b/checks/ps6136_head_width.go @@ -0,0 +1,168 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136ProjectorHeadWidth walks only source returns and closed value structs. +// Every executable alternative must carry the same selected head width. A +// configured projector type does not attest arbitrary opaque factory results. +func ps6136ProjectorHeadWidth(context *ps6125SSAContext, value ssa.Value, projectors map[*types.Named]*types.Var, model ps6125SSAReference, modelType *types.Named, head *types.Var, shape string, axis, budget int) bool { + if context == nil || value == nil || budget <= 0 { + return false + } + reference := context.reference(value) + if reference.value == nil { + return false + } + if reference.context != context || reference.value != value { + return ps6136ProjectorHeadWidth(reference.context, reference.value, projectors, model, modelType, head, shape, axis, budget-1) + } + if boxed, ok := value.(*ssa.MakeInterface); ok { + return ps6136ProjectorHeadWidth(context, boxed.X, projectors, model, modelType, head, shape, axis, budget-1) + } + if call, ok := value.(*ssa.Call); ok { + child := ps6136ZeroSourceFlow(ps6136Call(context, call)) + if child == nil { + return false + } + found := false + for _, block := range child.flow.function.Blocks { + if !child.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + ret, ok := instruction.(*ssa.Return) + if !ok { + continue + } + if len(ret.Results) != 1 || !ps6136ProjectorHeadWidth(child, ret.Results[0], projectors, model, modelType, head, shape, axis, budget-1) { + return false + } + found = true + } + } + return found + } + if phi, ok := value.(*ssa.Phi); ok { + found := false + for index, predecessor := range phi.Block().Preds { + if !context.flow.edges[ps6125SSAEdge{from: predecessor, to: phi.Block()}] { + continue + } + if !ps6136ProjectorHeadWidth(context, phi.Edges[index], projectors, model, modelType, head, shape, axis, budget-1) { + return false + } + found = true + } + return found + } + named, ok := types.Unalias(value.Type()).(*types.Named) + if !ok || projectors[named] == nil { + return false + } + structure, ok := named.Underlying().(*types.Struct) + if !ok { + return false + } + for index := 0; index < structure.NumFields(); index++ { + if structure.Field(index) != projectors[named] { + continue + } + width := context.structField(value, index) + return width.value != nil && ps6136HeadWidth(width.context, width.value, model, modelType, head, shape, axis, budget-1) + } + return false +} + +// ps6136HeadWidth proves only source flow from the selected model's exact head +// tensor through Shape()[axis] to a projector width. Equality with cfg.Vocab, +// valid tensor shape and native projector semantics require separate reviewed +// facts. Invariance and complete owner observations are independent gates. +func ps6136HeadWidth(context *ps6125SSAContext, value ssa.Value, model ps6125SSAReference, modelType *types.Named, head *types.Var, shape string, axis int, budget int) bool { + if context == nil || value == nil || model.value == nil || modelType == nil || head == nil || axis < 0 || budget <= 0 { + return false + } + reference := context.reference(value) + if reference.value == nil { + return false + } + if reference.context != context || reference.value != value { + return ps6136HeadWidth(reference.context, reference.value, model, modelType, head, shape, axis, budget-1) + } + if call, ok := value.(*ssa.Call); ok { + child := ps6136Call(context, call) + if child == nil { + return false + } + found := false + for _, block := range child.flow.function.Blocks { + if !child.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + ret, ok := instruction.(*ssa.Return) + if !ok { + continue + } + if len(ret.Results) != 1 || !ps6136HeadWidth(child, ret.Results[0], model, modelType, head, shape, axis, budget-1) { + return false + } + found = true + } + } + return found + } + if phi, ok := value.(*ssa.Phi); ok { + found := false + for index, predecessor := range phi.Block().Preds { + if !context.flow.edges[ps6125SSAEdge{from: predecessor, to: phi.Block()}] { + continue + } + if !ps6136HeadWidth(context, phi.Edges[index], model, modelType, head, shape, axis, budget-1) { + return false + } + found = true + } + return found + } + load, ok := value.(*ssa.UnOp) + if !ok || load.Op != token.MUL || !types.Identical(value.Type(), types.Typ[types.Int]) { + return false + } + index, ok := load.X.(*ssa.IndexAddr) + if !ok { + return false + } + integer := context.scalar(index.Index) + if integer.state != ps6125Integer || !ps6125SameExtent(integer.extent, ps6125ConstantExtent(int64(axis))) { + return false + } + base := context.reference(index.X) + call, ok := base.value.(*ssa.Call) + if !ok || call.Call.StaticCallee() == nil || len(call.Call.Args) != 1 { + return false + } + function, ok := call.Call.StaticCallee().Object().(*types.Func) + if !ok || ps6090FunctionID(function) != shape { + return false + } + signature, ok := function.Type().(*types.Signature) + if !ok || signature.Recv() == nil || signature.Variadic() || signature.TypeParams().Len() != 0 || signature.RecvTypeParams().Len() != 0 || signature.Params().Len() != 0 || signature.Results().Len() != 1 { + return false + } + slice, ok := call.Type().Underlying().(*types.Slice) + if !ok || !types.Identical(slice.Elem(), types.Typ[types.Int]) { + return false + } + receiver := base.context.reference(call.Call.Args[0]) + if receiver.value == nil { + return false + } + paths := ps6125AccessPaths{flow: receiver.context.flow} + path := paths.resolve(receiver.value) + return path.known && len(path.access.fields) == 1 && path.access.fields[0] == head && ps6136AccessOwnerRoot(receiver.context, path, modelType) == model +} diff --git a/checks/ps6136_head_width_test.go b/checks/ps6136_head_width_test.go new file mode 100644 index 00000000..40546bf7 --- /dev/null +++ b/checks/ps6136_head_width_test.go @@ -0,0 +1,46 @@ +package checks + +import ( + "go/types" + "testing" +) + +func TestPS6136HeadWidthSourceIdentity(t *testing.T) { + t.Parallel() + pkg := ps6125TestSSA(t, `package extent +type tensor struct{} +func(*tensor)Shape()[]int{return nil} +func(*tensor)Other()[]int{return nil} +type model struct{head,other *tensor} +func width(t *tensor)int{return t.Shape()[1]} +func leaf(int){} +func work(m,sibling *model){ + leaf(width(m.head));leaf(m.head.Shape()[1]);leaf(m.head.Shape()[0]) + leaf(sibling.head.Shape()[1]);leaf(m.other.Shape()[1]);leaf(m.head.Other()[1]) +} +`) + entry := pkg.Func("work") + context := ps6125NewSSAContext(entry, nil, nil, 128) + modelType := pkg.Pkg.Scope().Lookup("model").Type().(*types.Named) + head := modelType.Underlying().(*types.Struct).Field(0) + tensorType := pkg.Pkg.Scope().Lookup("tensor").Type().(*types.Named) + method, _, _ := types.LookupFieldOrMethod(tensorType, true, pkg.Pkg, "Shape") + model := context.reference(entry.Params[0]) + count := 0 + for _, call := range ps6125ContextCalls(context) { + if call.Call.StaticCallee() != pkg.Func("leaf") { + continue + } + got := ps6136HeadWidth(context, call.Call.Args[0], model, modelType, head, ps6090FunctionID(method.(*types.Func)), 1, 64) + if got != (count < 2) { + t.Fatalf("head witness %d = %v", count, got) + } + if ps6136HeadWidth(context, call.Call.Args[0], model, modelType, head, ps6090FunctionID(method.(*types.Func)), 1, 0) { + t.Fatal("exhausted proof budget accepted") + } + count++ + } + if count != 6 { + t.Fatalf("leaf count %d", count) + } +} diff --git a/checks/ps6136_host_prefix.go b/checks/ps6136_host_prefix.go new file mode 100644 index 00000000..27a419f4 --- /dev/null +++ b/checks/ps6136_host_prefix.go @@ -0,0 +1,152 @@ +package checks + +import ( + "go/ast" + "go/constant" + "go/token" + "go/types" + + "golang.org/x/tools/go/analysis" +) + +// ps6136HostPrefix distinguishes an unbounded physical ToHost transfer from +// its source-visible first-width logical consumption. It accepts only the +// actual closed Tensor.Storage().F32() chain and canonical prefix copy; extra +// tensor/slice observations, aliases or full consumers are unknown, not small. +func ps6136HostPrefix(pass *analysis.Pass, transfer *ast.CallExpr, receiver types.Object, width *types.Var, storage, f32 *types.Func) bool { + if pass == nil || transfer == nil || receiver == nil || width == nil || storage == nil || f32 == nil { + return false + } + parents := make(map[ast.Node]ast.Node) + for _, file := range pass.Files { + var stack []ast.Node + ast.Inspect(file, func(node ast.Node) bool { + if node == nil { + stack = stack[:len(stack)-1] + return true + } + if len(stack) > 0 { + parents[node] = stack[len(stack)-1] + } + stack = append(stack, node) + return true + }) + } + assignment, ok := parents[transfer].(*ast.AssignStmt) + if !ok || assignment.Tok != token.DEFINE || len(assignment.Rhs) != 1 || len(assignment.Lhs) != 2 { + return false + } + tensor, ok := assignment.Lhs[0].(*ast.Ident) + if !ok || pass.TypesInfo.Defs[tensor] == nil { + return false + } + tensorObject := pass.TypesInfo.Defs[tensor] + var tensorUse *ast.Ident + for identifier, object := range pass.TypesInfo.Uses { + if object != tensorObject { + continue + } + if tensorUse != nil { + return false + } + tensorUse = identifier + } + selector, ok := parents[tensorUse].(*ast.SelectorExpr) + if !ok || selector.X != tensorUse || !ps6136SelectedMethod(pass, selector, storage) { + return false + } + storageCall, ok := parents[selector].(*ast.CallExpr) + if !ok || storageCall.Fun != selector || len(storageCall.Args) != 0 { + return false + } + f32Selector, ok := parents[storageCall].(*ast.SelectorExpr) + if !ok || f32Selector.X != storageCall || !ps6136SelectedMethod(pass, f32Selector, f32) { + return false + } + f32Call, ok := parents[f32Selector].(*ast.CallExpr) + if !ok || f32Call.Fun != f32Selector || len(f32Call.Args) != 0 || !types.Identical(pass.TypesInfo.TypeOf(f32Call), types.NewSlice(types.Typ[types.Float32])) { + return false + } + sliceAssignment, ok := parents[f32Call].(*ast.AssignStmt) + if !ok || sliceAssignment.Tok != token.DEFINE || len(sliceAssignment.Lhs) != 1 || len(sliceAssignment.Rhs) != 1 { + return false + } + slice, ok := sliceAssignment.Lhs[0].(*ast.Ident) + if !ok || pass.TypesInfo.Defs[slice] == nil { + return false + } + var sliceUse *ast.Ident + for identifier, object := range pass.TypesInfo.Uses { + if object == pass.TypesInfo.Defs[slice] { + if sliceUse != nil { + return false + } + sliceUse = identifier + } + } + index, ok := parents[sliceUse].(*ast.IndexExpr) + if !ok || index.X != sliceUse { + return false + } + var loop *ast.ForStmt + for node := ast.Node(index); node != nil; node = parents[node] { + if candidate, ok := node.(*ast.ForStmt); ok { + loop = candidate + break + } + } + if loop == nil || len(loop.Body.List) != 1 { + return false + } + init, ok := loop.Init.(*ast.AssignStmt) + if !ok || init.Tok != token.DEFINE || len(init.Lhs) != 1 || len(init.Rhs) != 1 { + return false + } + i, ok := init.Lhs[0].(*ast.Ident) + if !ok || pass.TypesInfo.Defs[i] == nil { + return false + } + zero := pass.TypesInfo.Types[init.Rhs[0]].Value + if zero == nil || zero.Kind() != constant.Int || constant.Sign(zero) != 0 { + return false + } + condition, ok := loop.Cond.(*ast.BinaryExpr) + if !ok || condition.Op != token.LSS || !ps6136ObjectExpr(pass, condition.X, pass.TypesInfo.Defs[i]) || !ps6136DirectField(pass, condition.Y, receiver, width) { + return false + } + post, ok := loop.Post.(*ast.IncDecStmt) + if !ok || post.Tok != token.INC || !ps6136ObjectExpr(pass, post.X, pass.TypesInfo.Defs[i]) || !ps6136ObjectExpr(pass, index.Index, pass.TypesInfo.Defs[i]) { + return false + } + copy, ok := loop.Body.List[0].(*ast.AssignStmt) + if !ok || copy.Tok != token.ASSIGN || len(copy.Lhs) != 1 || len(copy.Rhs) != 1 { + return false + } + destination, ok := copy.Lhs[0].(*ast.IndexExpr) + if !ok || !ps6136ObjectExpr(pass, destination.Index, pass.TypesInfo.Defs[i]) { + return false + } + // A call/container expression on the assignment's left side may mutate + // i before the source index is evaluated. The actual owner copies to a + // bare []float64 variable; no effectful destination grammar is approved. + destinationBuffer, ok := destination.X.(*ast.Ident) + if !ok || pass.TypesInfo.Uses[destinationBuffer] == nil || !types.Identical(pass.TypesInfo.TypeOf(destinationBuffer), types.NewSlice(types.Typ[types.Float64])) { + return false + } + conversion, ok := copy.Rhs[0].(*ast.CallExpr) + if !ok || len(conversion.Args) != 1 || conversion.Args[0] != index || !types.Identical(pass.TypesInfo.TypeOf(conversion.Fun), types.Typ[types.Float64]) { + return false + } + allowedIteratorUses := map[ast.Node]bool{condition.X: true, post.X: true, index.Index: true, destination.Index: true} + for identifier, object := range pass.TypesInfo.Uses { + if object == pass.TypesInfo.Defs[i] && !allowedIteratorUses[identifier] { + return false + } + } + return true +} + +func ps6136SelectedMethod(pass *analysis.Pass, selector *ast.SelectorExpr, method *types.Func) bool { + selection := pass.TypesInfo.Selections[selector] + return selection != nil && selection.Kind() == types.MethodVal && selection.Obj() == method +} diff --git a/checks/ps6136_host_prefix_test.go b/checks/ps6136_host_prefix_test.go new file mode 100644 index 00000000..e55c0d47 --- /dev/null +++ b/checks/ps6136_host_prefix_test.go @@ -0,0 +1,118 @@ +package checks + +import ( + "go/ast" + "go/parser" + "go/token" + "go/types" + "testing" + + "golang.org/x/tools/go/analysis" +) + +func TestPS6136AuthenticCapacityTransferPrefix(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + t.Run(revision, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, revision, false, true) + if err != nil { + t.Fatal(err) + } + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg} + owner := fixture.pkg.Scope().Lookup("Decoder").Type().(*types.Named) + width, _, _ := types.LookupFieldOrMethod(owner, true, fixture.pkg, "v") + var tensorPackage *types.Package + for _, imported := range fixture.pkg.Imports() { + if imported.Path() == "github.com/jxsl13/goai/tensor" { + tensorPackage = imported + } + } + storage, _, _ := types.LookupFieldOrMethod(types.NewPointer(tensorPackage.Scope().Lookup("Tensor").Type()), true, fixture.pkg, "Storage") + f32, _, _ := types.LookupFieldOrMethod(types.NewPointer(tensorPackage.Scope().Lookup("Storage").Type()), true, fixture.pkg, "F32") + found := 0 + for _, file := range fixture.files { + for _, declaration := range file.Decls { + function, ok := declaration.(*ast.FuncDecl) + if !ok || function.Recv == nil || function.Name.Name != "Generate" || !types.Identical(fixture.info.TypeOf(function.Recv.List[0].Type), types.NewPointer(owner)) { + continue + } + ast.Inspect(function.Body, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if !ok { + return true + } + selector, ok := call.Fun.(*ast.SelectorExpr) + if !ok || selector.Sel.Name != "ToHost" { + return true + } + found++ + if !ps6136HostPrefix(pass, call, fixture.info.Defs[function.Recv.List[0].Names[0]], width.(*types.Var), storage.(*types.Func), f32.(*types.Func)) { + t.Fatal("authentic physical-capacity transfer's first-vocab logical consumption unproved") + } + return true + }) + } + } + if found != 1 { + t.Fatalf("authentic capacity transfer count %d", found) + } + }) + } +} + +func TestPS6136HostPrefixBoundaries(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name, body string + valid bool + }{ + {"prefix", `for i:=0;i= 0 { + if role >= len(arguments) { + return ps6125Extent{}, false + } + value = arguments[role] + } else if role != -1 { + return ps6125Extent{}, false + } + if !ps6136BufferOrigin(context, value, facts.owner, facts.ownerType, workspace, slot, 128) { + return ps6125Extent{}, false + } + } + integer := func(index int) ps6125Extent { + if index < 0 || index >= len(arguments) || !types.Identical(signature.Params().At(index).Type(), types.Typ[types.Int]) { + return ps6125Extent{} + } + return facts.extent(context, arguments[index]) + } + switch leaf.Kind { + case "projection": + if projector == nil { + return ps6125Extent{}, false + } + reference := context.reference(receiver) + if reference.value == nil { + return ps6125Extent{}, false + } + paths := ps6125AccessPaths{flow: reference.context.flow} + path := paths.resolve(reference.value) + if !path.known || len(path.access.fields) != 1 || path.access.fields[0] != projector || ps6136AccessOwnerRoot(reference.context, path, facts.ownerType) != facts.owner { + return ps6125Extent{}, false + } + return ps6125MultiplyExtents(integer(leaf.RowsArgument), facts.immutable[width]), true + case "row-width": + return ps6125MultiplyExtents(integer(leaf.RowsArgument), integer(leaf.WidthArgument)), true + case "prefix": + return integer(leaf.WidthArgument), true + case "download": + index := leaf.HostArgument + if index < 0 || index >= len(arguments) || !types.Identical(signature.Params().At(index).Type(), types.NewSlice(types.Typ[types.Float32])) { + return ps6125Extent{}, false + } + return facts.length(context, arguments[index]), true + case "capacity-transfer": + return ps6125Extent{}, true // Physical capacity-wide, separately prove logical prefix. + } + return ps6125Extent{}, false +} diff --git a/checks/ps6136_leaf_extent_test.go b/checks/ps6136_leaf_extent_test.go new file mode 100644 index 00000000..725218ec --- /dev/null +++ b/checks/ps6136_leaf_extent_test.go @@ -0,0 +1,114 @@ +package checks + +import ( + "go/types" + "testing" + + "github.com/jxsl13/perfscan/config" + "golang.org/x/tools/go/ssa" +) + +func TestPS6136AuthenticProjectionAndDownloadExtents(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, "before", true, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + for _, ownerName := range []string{"GPTDecoder", "Decoder"} { + ownerType := pkg.Pkg.Scope().Lookup(ownerName).Type().(*types.Named) + var step *ssa.Function + method, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, "Step") + step = pkg.Prog.FuncValue(method.(*types.Func)) + context := ps6125NewSSAContext(step, nil, nil, 256) + owner := context.reference(step.Params[0]) + workspace, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, "logits") + width, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, "v") + projectorName := "head" + if ownerName == "Decoder" { + projectorName = "out" + } + projector, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, projectorName) + slot := pkg.Pkg.Scope().Lookup("bufSlot").Type().Underlying().(*types.Struct).Field(0) + var pool ps6125ExtentPool + widthExtent := ps6125SymbolicExtent(pool.identity(step.Params[0].Object(), []*types.Var{width.(*types.Var)}, false)) + facts := &ps6136Extents{owner: owner, ownerType: ownerType, immutable: map[*types.Var]ps6125Extent{width.(*types.Var): widthExtent}, remaining: 2048} + count, projections := 0, 0 + provedLeaves := make(map[*ssa.Call]bool) + reads := make(map[ssa.Value]bool) + remaining := 4096 + if !ps6136WalkConsumerCalls(context, func(current *ps6125SSAContext, call *ssa.Call) bool { + for _, block := range current.flow.function.Blocks { + if !current.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + load, ok := instruction.(*ssa.UnOp) + if !ok { + continue + } + address, ok := load.X.(*ssa.FieldAddr) + if !ok { + continue + } + pointer, ok := address.X.Type().Underlying().(*types.Pointer) + if !ok || !types.Identical(pointer.Elem(), ownerType) { + continue + } + if ownerType.Underlying().(*types.Struct).Field(address.Field) == workspace { + reads[load] = true + } + } + } + if !call.Call.IsInvoke() { + return false + } + leaf := config.OutputWorkspaceLeaf{Method: ps6090FunctionID(call.Call.Method), Kind: "download", BufferArguments: []int{-1}, RowsArgument: -1, WidthArgument: -1, HostArgument: 0, SemanticsReviewed: true} + switch call.Call.Method.Name() { + case "DownloadF32": + count++ + case "record": + leaf.Kind, leaf.BufferArguments, leaf.RowsArgument, leaf.HostArgument = "projection", []int{2}, 3, -1 + // Other layers also invoke record; only the exact source owner + // head/out receiver may establish output-workspace extent. + if _, ok := ps6136LeafExtent(current, call, &leaf, facts, workspace.(*types.Var), slot, projector.(*types.Var), width.(*types.Var)); !ok { + return false + } + projections++ + case "AddBias": + leaf.Kind, leaf.BufferArguments, leaf.RowsArgument, leaf.WidthArgument, leaf.HostArgument = "row-width", []int{0, 2}, 3, 4, -1 + if _, ok := ps6136LeafExtent(current, call, &leaf, facts, workspace.(*types.Var), slot, projector.(*types.Var), width.(*types.Var)); !ok { + return false + } + default: + return false + } + got, ok := ps6136LeafExtent(current, call, &leaf, facts, workspace.(*types.Var), slot, projector.(*types.Var), width.(*types.Var)) + if !ok || !ps6125SameExtent(got, widthExtent) { + t.Fatal("actual one-row source leaf extent or workspace identity missing") + } + provedLeaves[call] = true + return true + }, &remaining) { + t.Fatal("source consumer traversal exhausted") + } + if count != 1 || projections != 1 { + t.Fatalf("%s source downloads %d projections %d", ownerName, count, projections) + } + expectedReads := 2 + if ownerName == "Decoder" { + expectedReads = 4 + } // optional output-bias input/output + if len(reads) != expectedReads { + t.Fatalf("%s workspace slot reads %d", ownerName, len(reads)) + } + for read := range reads { + if !ps6136WorkspaceReadClosed(read, slot, provedLeaves, 128) { + t.Fatalf("%s unclosed authentic read %s", ownerName, read) + } + if ps6136WorkspaceReadClosed(read, slot, nil, 128) { + t.Fatal("unproved consumer accepted") + } + } + } +} diff --git a/checks/ps6136_leaf_recorder.go b/checks/ps6136_leaf_recorder.go new file mode 100644 index 00000000..4788ccdc --- /dev/null +++ b/checks/ps6136_leaf_recorder.go @@ -0,0 +1,80 @@ +package checks + +import ( + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// leafRecorderOrigin joins the actual output call's recorder to this same +// published owner's source factory field. NativeBindings independently closes +// both fields and their only permitted profile override/native implementations. +// A matching interface method/signature alone never supplies this provenance. +func (selection *ps6136Selection) leafRecorderOrigin(context *ps6125SSAContext, call *ssa.Call, kind string, owner ps6125SSAReference) bool { + if selection == nil || context == nil || call == nil || selection.ops == nil || selection.primaryRecorder == nil { + return false + } + receiver := call.Call.Value + arguments := call.Call.Args + if !call.Call.IsInvoke() { + if len(arguments) == 0 { + return false + } + receiver, arguments = arguments[0], arguments[1:] + } + if kind == "projection" { + // The separately proved F32 lowering binds source parameter0 to the + // configured recorder projection receiver, not another same-typed field. + if len(arguments) == 0 { + return false + } + receiver = arguments[0] + } + reference := context.reference(receiver) + extraction, ok := reference.value.(*ssa.Extract) + if !ok || extraction.Index != 0 { + return false + } + factory, ok := extraction.Tuple.(*ssa.Call) + if !ok || factory.Call.IsInvoke() { + return false + } + signature := factory.Call.Signature() + if signature.Variadic() || signature.TypeParams().Len() != 0 || signature.Params().Len() != 0 || len(factory.Call.Args) != 0 || signature.Results().Len() != 2 || !types.Identical(signature.Results().At(0).Type(), receiver.Type()) || !types.Identical(signature.Results().At(1).Type(), types.Universe.Lookup("error").Type()) { + return false + } + remaining := 128 + // Keep the actual read/phi: callable normalization can replace the profile + // field's source read with its known closure, losing the owner-field join. + factoryReference := ps6125SSAReference{context: reference.context, value: factory.Call.Value} + return selection.recorderFactoryOrigin(factoryReference, owner, make(map[ps6125SSAReference]bool), &remaining) +} + +func (selection *ps6136Selection) recorderFactoryOrigin(reference ps6125SSAReference, owner ps6125SSAReference, active map[ps6125SSAReference]bool, remaining *int) bool { + if reference.value == nil || reference.context == nil || *remaining <= 0 || active[reference] { + return false + } + *remaining-- + active[reference] = true + defer delete(active, reference) + paths := ps6125AccessPaths{flow: reference.context.flow} + path := paths.resolve(reference.value) + if path.known && len(path.access.fields) == 2 && path.access.fields[0] == selection.ops && (path.access.fields[1] == selection.primaryRecorder || selection.secondaryRecorder != nil && path.access.fields[1] == selection.secondaryRecorder) { + return ps6136OwnerRoot(reference.context, path.access.root, selection.ownerType) == owner + } + phi, ok := reference.value.(*ssa.Phi) + if !ok || len(phi.Edges) != len(phi.Block().Preds) { + return false + } + seen := false + for index, edge := range phi.Edges { + if !reference.context.flow.blocks[phi.Block().Preds[index]] { + continue + } + if !selection.recorderFactoryOrigin(ps6125SSAReference{context: reference.context, value: edge}, owner, active, remaining) { + return false + } + seen = true + } + return seen +} diff --git a/checks/ps6136_native_bindings.go b/checks/ps6136_native_bindings.go new file mode 100644 index 00000000..c44cb699 --- /dev/null +++ b/checks/ps6136_native_bindings.go @@ -0,0 +1,374 @@ +package checks + +import ( + "go/ast" + "go/types" + "slices" + + "github.com/jxsl13/perfscan/config" + "golang.org/x/tools/go/ssa" +) + +func (context *ps6136ContractContext) nativeBindings(selection *ps6136Selection, proof *ps6136ConstructorProof, contract *config.OutputWorkspaceContract) bool { + if selection == nil || proof == nil || contract == nil { + return false + } + base := context.named(contract.RecorderWrapperType) + native := context.named(contract.NativeRecorderType) + buffer := context.named(contract.BufferWrapperType) + projector := context.named(contract.ProjectorType) + if base == nil || native == nil || buffer == nil || projector == nil { + return false + } + nativeField := ps6136FieldVar(base, contract.RecorderNativeField) + bufferField := ps6136FieldVar(buffer, contract.BufferNativeField) + recorderField := ps6136FieldVar(selection.ops.Type(), contract.BackendRecorderField) + if nativeField == nil || bufferField == nil || recorderField == nil || !types.Identical(nativeField.Type(), types.NewPointer(native)) || !selection.recorderBinding(context.pass, recorderField, contract.NativeRecorderFactory, base, nativeField) { + return false + } + if len(proof.factory.flow.function.Params) != 1 { + return false + } + backend := ps6136FactoryResult(proof.factory, proof.factory.reference(proof.factory.flow.function.Params[0]), selection.allocator, selection.slot) + if backend == nil || !ps6136BoundBufferWrapper(proof.factory, backend, buffer, bufferField) { + return false + } + if contract.NativeBufferReleaseMethod != "" { + release := context.callable(contract.RetainedBufferReleaseMethod) + if release == nil || release.Type().(*types.Signature).Recv() == nil { + return false + } + object, _, _ := types.LookupFieldOrMethod(buffer, false, context.pass.Pkg, release.Name()) + nativeRelease, ok := object.(*types.Func) + if !ok || nativeRelease.Type().(*types.Signature).Recv() == nil || ps6090FunctionID(nativeRelease) != contract.NativeBufferReleaseMethod || nativeRelease.Type().(*types.Signature).Params().Len() != 0 || nativeRelease.Type().(*types.Signature).Results().Len() != 0 { + return false + } + } + bridge := context.callable(contract.BufferBridge) + if bridge == nil || !ps6136BufferBridge(context.pass, context.declarations[bridge], buffer, bufferField) { + return false + } + record := context.callable(contract.ProjectorRecordMethod) + projection := context.callable(contract.RecorderProjectionMethod) + if record == nil || projection == nil || record.Type().(*types.Signature).Recv() == nil || projection.Type().(*types.Signature).Recv() == nil || !types.Identical(record.Type().(*types.Signature).Recv().Type(), projector) { + return false + } + weight := ps6136FieldVar(projector, contract.ProjectorWeightField) + inner := ps6136FieldVar(projector, contract.ProjectorInnerField) + width := ps6136FieldVar(projector, contract.ProjectorWidthField) + if weight == nil || inner == nil || width == nil || !types.Identical(inner.Type(), types.Typ[types.Int]) || !types.Identical(width.Type(), types.Typ[types.Int]) || selection.projectorWidths[projector] != width || !ps6136ProjectionLowering(context.pass, context.declarations[record], contract.RecorderProjectionMethod, weight, inner, width) { + return false + } + object, _, _ := types.LookupFieldOrMethod(base, false, context.pass.Pkg, projection.Name()) + adapter, ok := object.(*types.Func) + if !ok || !ps6136ForwardedLeaf(context.pass, context.declarations[adapter], contract.NativeProjectionMethod, bridge, []int{0, 1, 2}) { + return false + } + adapters := []*types.Func{adapter} + var buffers []int + for index := range contract.Leaves { + leaf := &contract.Leaves[index] + method := context.callable(leaf.Method) + if method == nil { + return false + } + signature := method.Type().(*types.Signature) + if signature.Recv() == nil { + return false + } + capability, ok := signature.Recv().Type().Underlying().(*types.Interface) + if !ok { + return false + } + if leaf.Kind == "projection" { + if len(leaf.Implementations) != 1 || leaf.Implementations[0] != contract.ProjectorRecordMethod || !types.Implements(projector, capability) { + return false + } + continue + } + if leaf.Kind == "row-width" { + object, _, _ := types.LookupFieldOrMethod(base, false, context.pass.Pkg, method.Name()) + adapter, ok := object.(*types.Func) + if !ok || len(leaf.Implementations) != 1 { + return false + } + buffers = buffers[:0] + for index := 0; index < signature.Params().Len(); index++ { + if types.Identical(signature.Params().At(index).Type(), selection.slot.Type()) { + buffers = append(buffers, index) + } + } + if !ps6136ForwardedLeaf(context.pass, context.declarations[adapter], leaf.Implementations[0], bridge, buffers) { + return false + } + adapters = append(adapters, adapter) + continue + } + if !types.Implements(buffer, capability) { + // The exact selected factory type cannot enter this optional + // capability branch. Its source/logical inventory remains required. + if !leaf.SelectedCapabilityAbsenceAllowed { + return false + } + continue + } + object, _, _ := types.LookupFieldOrMethod(buffer, false, context.pass.Pkg, method.Name()) + implementation, ok := object.(*types.Func) + if !ok || len(leaf.Implementations) != 1 || ps6090FunctionID(implementation) != leaf.Implementations[0] { + return false + } + // Promotion through the exact one-field allocator wrapper must reach + // the native buffer, not a new wrapper-owned override. + pointer, ok := bufferField.Type().(*types.Pointer) + if !ok || !bufferField.Embedded() || implementation.Type().(*types.Signature).Recv() == nil || !types.Identical(implementation.Type().(*types.Signature).Recv().Type(), pointer) { + return false + } + } + return context.recorderStateClosed(selection, contract, recorderField, base, nativeField, adapters) +} + +func ps6136LiteralFieldInitializers(context *ps6136ContractContext, fields map[*types.Var]bool) map[ast.Node]bool { + proved := make(map[ast.Node]bool) + for _, file := range context.pass.Files { + ast.Inspect(file, func(node ast.Node) bool { + if pair, ok := node.(*ast.KeyValueExpr); ok { + if key, ok := pair.Key.(*ast.Ident); ok { + if field, ok := context.pass.TypesInfo.Uses[key].(*types.Var); ok && fields[field] { + proved[pair] = true + } + } + } + return true + }) + } + return proved +} + +func (context *ps6136ContractContext) recorderStateClosed(selection *ps6136Selection, contract *config.OutputWorkspaceContract, primary *types.Var, base *types.Named, nativeField *types.Var, adapters []*types.Func) bool { + fields := map[*types.Var]bool{primary: true} + secondary := ps6136FieldVar(selection.ops.Type(), contract.SecondaryRecorderField) + if contract.SecondaryRecorderField != "" { + if secondary == nil { + return false + } + fields[secondary] = true + snapshot := ps6136InitialOwnerField(selection.owner, selection.ops) + index := -1 + structure := selection.ops.Type().Underlying().(*types.Struct) + for candidate := 0; candidate < structure.NumFields(); candidate++ { + if structure.Field(candidate) == secondary { + index = candidate + } + } + if snapshot.value == nil || index < 0 || !ps6136StructFieldZero(snapshot.context, snapshot.value, index, 512) { + return false + } + } + proved := ps6136LiteralFieldInitializers(context, fields) + for _, file := range context.pass.Files { + ast.Inspect(file, func(node ast.Node) bool { + selector, ok := node.(*ast.SelectorExpr) + if !ok { + return true + } + selectionInfo := context.pass.TypesInfo.Selections[selector] + if selectionInfo == nil { + return true + } + field, ok := selectionInfo.Obj().(*types.Var) + if !ok || !fields[field] { + return true + } + container, ok := selector.X.(*ast.SelectorExpr) + if !ok { + return true + } + containerInfo := context.pass.TypesInfo.Selections[container] + if containerInfo != nil && containerInfo.Kind() == types.FieldVal && containerInfo.Obj() != selection.ops && types.Identical(containerInfo.Type(), selection.ops.Type()) { + // A distinct owner's by-value backendOps field is different + // storage; it cannot mutate this selected owner's callback cell. + proved[selector] = true + } + return true + }) + } + stableSelected := true + remaining := 16384 + if !ps6136WalkConsumerCalls(selection.context, func(current *ps6125SSAContext, call *ssa.Call) bool { + for _, current := range []*ps6125SSAContext{current, ps6136Call(current, call)} { + if current == nil { + continue + } + for _, block := range current.flow.function.Blocks { + if !current.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + store, ok := instruction.(*ssa.Store) + if !ok { + continue + } + address, ok := store.Addr.(*ssa.FieldAddr) + if !ok { + continue + } + pointer, ok := address.X.Type().Underlying().(*types.Pointer) + if !ok { + continue + } + structure, ok := pointer.Elem().Underlying().(*types.Struct) + if !ok || !fields[structure.Field(address.Field)] { + continue + } + root := address.X + for { + parent, ok := root.(*ssa.FieldAddr) + if !ok { + break + } + root = parent.X + } + if types.Identical(root.Type(), types.NewPointer(selection.ownerType)) && ps6136OwnerRoot(current, root, selection.ownerType) == selection.owner { + stableSelected = false + } + } + } + } + return !stableSelected + }, &remaining) || !stableSelected { + return false + } + // Unselected constructors may specialize their own fresh instance, but + // cannot grant an exemption to calls on a published selected owner. + pkg := selection.context.flow.function.Pkg + writes := ps6136FreshConstructorWrites(pkg, selection.ownerType, fields) + if writes == nil { + return false + } + for _, file := range context.pass.Files { + ast.Inspect(file, func(node ast.Node) bool { + if selector, ok := node.(*ast.SelectorExpr); ok { + if selection := context.pass.TypesInfo.Selections[selector]; selection != nil { + if field, ok := selection.Obj().(*types.Var); ok && fields[field] { + for position := range writes.positions { + if selector.Pos() <= position && position < selector.End() { + proved[selector] = true + } + } + } + } + } + return true + }) + } + profileTargets := make(map[ast.Node]bool) + if contract.RecorderOverrideMethod != "" { + method := context.callable(contract.RecorderOverrideMethod) + wrapper := context.named(contract.RecorderOverrideWrapperType) + if method == nil || wrapper == nil || secondary == nil || !ps6136SameAdapterMethods(wrapper, base, adapters) { + return false + } + embedded := ps6136FieldVar(wrapper, contract.RecorderOverrideEmbedField) + metadata := ps6136FieldVar(wrapper, contract.RecorderOverrideMetadataField) + async := ps6136FieldVar(selection.ops.Type(), contract.AsyncRecorderField) + frame := ps6136ProfileOverrideFrame(context.pass, context.declarations[method], selection.ops, primary, secondary, async, contract.OwnerType+".Step", contract.DropPendingMethod, func(literal *ast.FuncLit, capacity, profile types.Object) bool { + return ps6136ProfileRecorderFactory(context.pass, literal, contract.NativeOverrideFactory, capacity, profile, wrapper, base, embedded, metadata, nativeField) + }) + if frame == nil { + return false + } + for node := range frame { + proved[node] = true + profileTargets[node] = true + } + } + for _, identity := range append(slices.Clone(contract.CommonMethods), contract.BulkMethod) { + method := context.callable(identity) + if method == nil { + return false + } + function := pkg.Prog.FuncValue(method) + if function == nil || len(function.Params) == 0 || !types.Identical(function.Params[0].Type(), types.NewPointer(selection.ownerType)) { + return false + } + root := ps6125NewSSAContext(function, nil, nil, 16384) + allowed := make(map[ast.Node]bool) + if identity == contract.RecorderOverrideMethod { + allowed = profileTargets + } + if !ps6136InstanceFieldWritesClosed(root, root.reference(function.Params[0]), selection.ownerType, fields, allowed) { + return false + } + if !ps6136InstanceFieldWritesClosed(root, root.reference(function.Params[0]), selection.ownerType, map[*types.Var]bool{selection.retained: true}, nil) { + return false + } + } + for field := range fields { + if !ps6136FieldEffects(context.pass, field, proved) { + return false + } + } + return true +} + +func ps6136InstanceFieldWritesClosed(root *ps6125SSAContext, owner ps6125SSAReference, ownerType *types.Named, fields map[*types.Var]bool, allowed map[ast.Node]bool) bool { + pending := []*ps6125SSAContext{root} + remaining := 16384 + for len(pending) != 0 { + current := pending[len(pending)-1] + pending = pending[:len(pending)-1] + remaining-- + if remaining < 0 { + return false + } + for _, block := range current.flow.function.Blocks { + if !current.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + if call, ok := instruction.(*ssa.Call); ok { + if child := ps6136Call(current, call); child != nil { + pending = append(pending, child) + } + } + store, ok := instruction.(*ssa.Store) + if !ok { + continue + } + address, ok := store.Addr.(*ssa.FieldAddr) + if !ok { + continue + } + pointer, ok := address.X.Type().Underlying().(*types.Pointer) + if !ok { + continue + } + structure, ok := pointer.Elem().Underlying().(*types.Struct) + if !ok || !fields[structure.Field(address.Field)] { + continue + } + base := address.X + for { + parent, ok := base.(*ssa.FieldAddr) + if !ok { + break + } + base = parent.X + } + if !types.Identical(base.Type(), types.NewPointer(ownerType)) || ps6136OwnerRoot(current, base, ownerType) != owner { + continue + } + proved := false + for target := range allowed { + if target.Pos() <= store.Pos() && store.Pos() < target.End() { + proved = true + } + } + if !proved { + return false + } + } + } + } + return true +} diff --git a/checks/ps6136_observation_invocation_test.go b/checks/ps6136_observation_invocation_test.go new file mode 100644 index 00000000..dd261a7d --- /dev/null +++ b/checks/ps6136_observation_invocation_test.go @@ -0,0 +1,138 @@ +package checks + +import ( + "go/ast" + "go/parser" + "go/token" + "go/types" + "testing" + + "golang.org/x/tools/go/ssa" +) + +func TestPS6136SharedHelperUnknownSecondInvocation(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwnersRewrite(t, "before", true, true, func(fs *token.FileSet, files []*ast.File) []*ast.File { + changed := 0 + for _, file := range files { + ast.Inspect(file, func(node ast.Node) bool { + if call, ok := node.(*ast.CallExpr); ok { + if selector, ok := call.Fun.(*ast.SelectorExpr); ok && selector.Sel.Name == "recordLogits" { + call.Args = append(call.Args, &ast.Ident{Name: "true"}) + } + } + return true + }) + } + for _, file := range files { + for _, declaration := range file.Decls { + function, ok := declaration.(*ast.FuncDecl) + if !ok { + continue + } + if function.Name.Name == "recordLogits" { + function.Type.Params.List = append(function.Type.Params.List, &ast.Field{Names: []*ast.Ident{{Name: "useKnown"}}, Type: &ast.Ident{Name: "bool"}}) + ast.Inspect(function.Body, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if !ok { + return true + } + selector, ok := call.Fun.(*ast.SelectorExpr) + if !ok { + return true + } + if selector.Sel.Name == "record" { + selector.X = &ast.Ident{Name: "projector"} + } + // Keep the bias extent independently bounded so rejection + // must come from the unproved projection binding, not an + // unrelated unknown AddBias extent. + if selector.Sel.Name == "AddBias" { + call.Args[3] = &ast.BasicLit{Kind: token.INT, Value: "1"} + } + return true + }) + prefix, err := parser.ParseFile(fs, "unknown-projector-adversary.go", "package llamagpu;func temporary(){projector:=d.out;if !useKnown{projector=unknownProjector()}}", parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + function.Body.List = append(prefix.Decls[0].(*ast.FuncDecl).Body.List, function.Body.List...) + } + if function.Name.Name != "encodeStep" { + continue + } + ast.Inspect(function.Body, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if !ok { + return true + } + selector, ok := call.Fun.(*ast.SelectorExpr) + if !ok || selector.Sel.Name != "recordLogits" { + return true + } + second, err := parser.ParseExpr("d.recordLogits(r,unknownRows(),unknownFlag())") + if err != nil { + t.Fatal(err) + } + original := *call + call.Fun = &ast.Ident{Name: "firstErr", NamePos: selector.Pos()} + call.Args = []ast.Expr{&original, second} + changed++ + return false + }) + } + } + if changed != 1 { + t.Fatalf("authentic encodeStep projection mutations %d", changed) + } + extra, err := parser.ParseFile(fs, "unknown-rows-adversary.go", "package llamagpu;var unknownRowsValue int;func unknownRows()int{return unknownRowsValue};var unknownProjectorValue linear;func unknownProjector()linear{return unknownProjectorValue};var unknownFlagValue bool;func unknownFlag()bool{return unknownFlagValue}", parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + return append(files, extra) + }) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + owner := pkg.Pkg.Scope().Lookup("Decoder").Type().(*types.Named) + field := func(name string) *types.Var { + object, _, _ := types.LookupFieldOrMethod(owner, true, pkg.Pkg, name) + return object.(*types.Var) + } + selection := &ps6136Selection{ownerType: owner, workspace: field("logits"), width: field("v"), projector: field("out"), slot: pkg.Pkg.Scope().Lookup("bufSlot").Type().Underlying().(*types.Struct).Field(0)} + selection.ops = field("ops") + selection.primaryRecorder = ps6136FieldVar(selection.ops.Type(), "newRecorder") + selection.secondaryRecorder = ps6136FieldVar(selection.ops.Type(), "newDecodeRecorder") + step, _, _ := types.LookupFieldOrMethod(owner, true, pkg.Pkg, "Step") + context := ps6125NewSSAContext(pkg.Prog.FuncValue(step.(*types.Func)), nil, nil, 16384) + if selection.consumerCalls(context, ps6136OwnerLeaves(pkg), false, nil) != nil { + t.Fatal("one proved invocation masked the same helper's unknown-row invocation") + } + var pool ps6125ExtentPool + width := ps6125SymbolicExtent(pool.identity(context.flow.function.Params[0].Object(), []*types.Var{selection.width}, false)) + facts := &ps6136Extents{owner: context.reference(context.flow.function.Params[0]), ownerType: owner, immutable: map[*types.Var]ps6125Extent{selection.width: width}, remaining: 16384} + valid, invalid := 0, 0 + var instruction *ssa.Call + remaining := 16384 + if !ps6136WalkConsumerCalls(context, func(current *ps6125SSAContext, call *ssa.Call) bool { + if current.flow.function.Name() != "recordLogits" || !call.Call.IsInvoke() || call.Call.Method.Name() != "record" { + return false + } + if instruction != nil && instruction != call { + t.Fatal("counterexample did not share one SSA leaf instruction") + } + instruction = call + if _, ok := ps6136LeafExtent(current, call, &ps6136OwnerLeaves(pkg)[0], facts, selection.workspace, selection.slot, selection.projector, selection.width); ok { + valid++ + } else { + invalid++ + } + return true + }, &remaining) { + t.Fatal("counterexample traversal exhausted") + } + if valid < 1 || invalid < 1 || valid != invalid { + t.Fatalf("same-instruction invocation bindings valid=%d invalid=%d", valid, invalid) + } +} diff --git a/checks/ps6136_observation_reads.go b/checks/ps6136_observation_reads.go new file mode 100644 index 00000000..dbd2c0cf --- /dev/null +++ b/checks/ps6136_observation_reads.go @@ -0,0 +1,99 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136WorkspaceReadClosed closes the complete SSA use graph of a workspace +// slot load against independently validated leaf calls. It grants no leaf +// semantics itself. Storing a slot/buffer alias, returning it, observing other +// slot members or passing it to any unproved call rejects the read. +func ps6136WorkspaceReadClosed(value ssa.Value, slot *types.Var, proved map[*ssa.Call]bool, budget int) bool { + if value == nil || slot == nil || budget <= 0 { + return false + } + seen := make(map[ssa.Value]bool) + var visit func(ssa.Value, bool) bool + visit = func(current ssa.Value, buffer bool) bool { + if seen[current] { + return true + } + seen[current] = true + budget-- + if budget < 0 || current.Referrers() == nil { + return false + } + used := false + for _, referrer := range *current.Referrers() { + switch use := referrer.(type) { + case *ssa.DebugRef: + continue + case *ssa.FieldAddr: + if buffer || use.X != current { + return false + } + pointer, ok := current.Type().Underlying().(*types.Pointer) + if !ok { + return false + } + structure, ok := pointer.Elem().Underlying().(*types.Struct) + if !ok || use.Field >= structure.NumFields() || structure.Field(use.Field) != slot { + return false + } + if use.Referrers() == nil || len(*use.Referrers()) == 0 { + return false + } + for _, reference := range *use.Referrers() { + if _, debug := reference.(*ssa.DebugRef); debug { + continue + } + load, ok := reference.(*ssa.UnOp) + if !ok || load.Op != token.MUL || load.X != use || !visit(load, true) { + return false + } + } + case *ssa.TypeAssert: + if !buffer || use.X != current { + return false + } + if use.CommaOk { + if use.Referrers() == nil { + return false + } + for _, reference := range *use.Referrers() { + if _, debug := reference.(*ssa.DebugRef); debug { + continue + } + extract, ok := reference.(*ssa.Extract) + if !ok || extract.Tuple != use || extract.Index > 1 { + return false + } + if extract.Index == 0 && !visit(extract, true) { + return false + } + // The assertion's boolean is a capability predicate, not + // an alias of its buffer. It cannot mutate buffer storage. + } + } else if !visit(use, true) { + return false + } + case *ssa.ChangeInterface: + if !buffer || use.X != current || !visit(use, true) { + return false + } + case *ssa.Call: + if !buffer || !proved[use] { + return false + } + default: + return false + } + used = true + } + return used + } + return visit(value, false) +} diff --git a/checks/ps6136_owner_adversaries_test.go b/checks/ps6136_owner_adversaries_test.go new file mode 100644 index 00000000..f6c4f36e --- /dev/null +++ b/checks/ps6136_owner_adversaries_test.go @@ -0,0 +1,174 @@ +package checks + +import ( + "github.com/jxsl13/perfscan/config" + "go/ast" + "go/parser" + "go/token" + "golang.org/x/tools/go/analysis" + "testing" +) + +func TestPS6136RegisteredOwnerSourceAdversaries(t *testing.T) { + t.Parallel() + for _, owner := range []string{"GPTDecoder", "Decoder"} { + for _, mutation := range []string{"release_early_return", "late_error_write", "late_allocation", "allocator_rebind", "constructor_workspace_release", "ledger_alias_release", "ledger_drop", "geometry_write", "whole_owner_escape", "full_logical_download", "wrong_recorder", "wrong_wrapper"} { + t.Run(owner+"/"+mutation, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwnersRewrite(t, "before", true, true, func(fs *token.FileSet, files []*ast.File) []*ast.File { + constructor := "newGPTDecoder" + if owner == "Decoder" { + constructor = "newDecoder" + } + parseStatements := func(body string) []ast.Stmt { + file, err := parser.ParseFile(fs, "adversary-"+mutation+".go", "package llamagpu;func temporary(){"+body+"}", parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + return file.Decls[0].(*ast.FuncDecl).Body.List + } + changed := 0 + for _, file := range files { + for _, declaration := range file.Decls { + function, ok := declaration.(*ast.FuncDecl) + if !ok { + continue + } + if function.Name.Name == constructor { + var body string + switch mutation { + case "late_error_write": + body = "err=nil" + case "late_allocation": + body = "_=mk(nil)" + case "allocator_rebind": + body = "d.ops.newBuffer=nil" + case "constructor_workspace_release": + body = "d.logits.b.Release()" + case "ledger_alias_release": + body = "d.all[len(d.all)-1].Release()" + case "whole_owner_escape": + body = "_=func(){_=d}" + } + if body != "" { + index := len(function.Body.List) - 1 + function.Body.List = append(function.Body.List[:index], append(parseStatements(body), function.Body.List[index:]...)...) + changed++ + } + } + receiver := "" + if function.Recv != nil { + if pointer, ok := function.Recv.List[0].Type.(*ast.StarExpr); ok { + if identifier, ok := pointer.X.(*ast.Ident); ok { + receiver = identifier.Name + } + } + } + if receiver == owner && function.Name.Name == "Release" && mutation == "release_early_return" { + function.Body.List = append(parseStatements("if d.v>0{return}"), function.Body.List...) + changed++ + } + if receiver == owner && function.Name.Name == "Step" { + body := "" + switch mutation { + case "ledger_drop": + body = "d.all=nil" + case "geometry_write": + body = "d.v++" + } + if body != "" { + function.Body.List = append(parseStatements(body), function.Body.List...) + changed++ + } + if mutation == "full_logical_download" { + ast.Inspect(function.Body, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if !ok { + return true + } + identifier, ok := call.Fun.(*ast.Ident) + if !ok || identifier.Name != "make" || len(call.Args) != 2 { + return true + } + selector, ok := call.Args[1].(*ast.SelectorExpr) + if !ok || selector.Sel.Name != "v" { + return true + } + call.Args[1] = &ast.BinaryExpr{X: &ast.SelectorExpr{X: &ast.Ident{Name: "d"}, Sel: &ast.Ident{Name: "maxLen"}}, Op: token.MUL, Y: selector} + changed++ + return true + }) + } + } + entry := "NewGPT" + if owner == "Decoder" { + entry = "New" + } + if function.Name.Name == entry && (mutation == "wrong_recorder" || mutation == "wrong_wrapper") { + ast.Inspect(function.Body, func(node ast.Node) bool { + pair, ok := node.(*ast.KeyValueExpr) + if !ok { + return true + } + key, ok := pair.Key.(*ast.Ident) + if !ok { + return true + } + if mutation == "wrong_recorder" && key.Name == "newRecorder" { + literal := pair.Value.(*ast.FuncLit) + call := literal.Body.List[0].(*ast.AssignStmt).Rhs[0].(*ast.CallExpr) + call.Fun.(*ast.SelectorExpr).Sel.Name = "NewConcurrentRecorder" + changed++ + } + if mutation == "wrong_wrapper" && key.Name == "newBuffer" { + literal := pair.Value.(*ast.FuncLit) + value := literal.Body.List[2].(*ast.ReturnStmt).Results[0].(*ast.CompositeLit) + value.Type = &ast.Ident{Name: "otherBufferWrapper"} + changed++ + } + return true + }) + } + } + } + if mutation == "wrong_wrapper" { + extra, err := parser.ParseFile(fs, "other-wrapper.go", "package llamagpu;type otherBufferWrapper struct{mBuf}", parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + files = append(files, extra) + // The nested original native value preserves the ABI while changing wrapper identity. + for _, file := range files { + ast.Inspect(file, func(node ast.Node) bool { + value, ok := node.(*ast.CompositeLit) + if ok { + if identifier, ok := value.Type.(*ast.Ident); ok && identifier.Name == "otherBufferWrapper" { + value.Elts = []ast.Expr{&ast.CompositeLit{Type: &ast.Ident{Name: "mBuf"}, Elts: value.Elts}} + } + } + return true + }) + } + } + if changed != 1 { + t.Fatalf("owner adversary rewrites %d want 1", changed) + } + return files + }) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + c := ps6136CompleteOwnerContract(pkg, owner) + var diagnostics []analysis.Diagnostic + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg, Report: func(d analysis.Diagnostic) { diagnostics = append(diagnostics, d) }} + if _, err := runPS6136WithContracts(pass, []config.OutputWorkspaceContract{c}); err != nil { + t.Fatal(err) + } + if len(diagnostics) != 0 { + t.Fatalf("unsafe owner-source candidate %d", len(diagnostics)) + } + }) + } + } +} diff --git a/checks/ps6136_owner_root.go b/checks/ps6136_owner_root.go new file mode 100644 index 00000000..8f288f46 --- /dev/null +++ b/checks/ps6136_owner_root.go @@ -0,0 +1,189 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136OwnerRoot resolves an exact owner pointer through a closed captured +// pointer cell. Unlike PS6125 callable cells, this is an identity-only proof: +// it establishes neither immutable owner memory nor device-resource lifetime. +func ps6136OwnerRoot(context *ps6125SSAContext, value ssa.Value, owner *types.Named) ps6125SSAReference { + if context == nil || value == nil || owner == nil || !types.Identical(value.Type(), types.NewPointer(owner)) { + return ps6125SSAReference{} + } + reference := context.reference(value) + seen := make(map[ps6125SSAReference]bool) + for reference.value != nil && !seen[reference] { + seen[reference] = true + call, isCall := reference.value.(*ssa.Call) + resultIndex := 0 + if extraction, ok := reference.value.(*ssa.Extract); ok { + call, isCall = extraction.Tuple.(*ssa.Call) + resultIndex = extraction.Index + } + if isCall { + child := ps6136Call(reference.context, call) + if child == nil { + break + } + var returnedRoot ps6125SSAReference + valid := true + for _, block := range child.flow.function.Blocks { + if !child.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + returned, ok := instruction.(*ssa.Return) + if !ok { + continue + } + if resultIndex >= len(returned.Results) { + valid = false + continue + } + value := returned.Results[resultIndex] + // Nil is not a competing owner identity. The caller must still + // establish successful construction before resource use. + if constant, ok := value.(*ssa.Const); ok && constant.IsNil() { + continue + } + root := ps6136OwnerRoot(child, value, owner) + _, fresh := root.value.(*ssa.Alloc) + if !fresh || returnedRoot.value != nil && returnedRoot != root { + valid = false + } + returnedRoot = root + } + } + if !valid || returnedRoot.value == nil { + break + } + reference = returnedRoot + continue + } + switch reference.value.(type) { + case *ssa.Parameter, *ssa.Alloc: + if types.Identical(reference.value.Type(), types.NewPointer(owner)) { + return reference + } + } + load, ok := reference.value.(*ssa.UnOp) + if !ok || load.Op != token.MUL { + break + } + cell := reference.context.reference(load.X) + allocation, ok := cell.value.(*ssa.Alloc) + if !ok || !types.Identical(allocation.Type(), types.NewPointer(types.NewPointer(owner))) { + break + } + store := ps6136ClosedOwnerCell(cell.context, allocation) + if store == nil { + break + } + var point ssa.Instruction = load + current := reference.context + for current != cell.context && current.parent != nil { + point = current.site + current = current.parent + } + if current != cell.context { + // A returned closure's creating function is not on the invocation + // stack. Require this exact closed cell among its source captures + // and initialization before creation (stronger than before return). + point = ps6136ReturnedCapturePoint(reference.context, cell) + } + if point == nil || !cell.context.flow.instructionDominates(store, point) { + break + } + reference = cell.context.reference(store.Val) + } + return ps6125SSAReference{} +} + +func ps6136ReturnedCapturePoint(context *ps6125SSAContext, cell ps6125SSAReference) ssa.Instruction { + for current := context; current != nil && current.parent != nil; current = current.parent { + if current.site == nil { + continue + } + origin := ps6136ReturnedCallable(current.parent.reference(current.site.Call.Value), make(map[ps6125SSAReference]bool)) + closure, ok := origin.value.(*ssa.MakeClosure) + if !ok || origin.context != cell.context { + continue + } + for _, capture := range closure.Bindings { + if origin.context.reference(capture) == cell { + return closure + } + } + } + return nil +} + +func ps6136ClosedOwnerCell(context *ps6125SSAContext, allocation *ssa.Alloc) *ssa.Store { + if context == nil || allocation.Parent() != context.flow.function || !context.flow.blocks[allocation.Block()] { + return nil + } + pending := []ssa.Value{allocation} + seen := make(map[ssa.Value]bool) + var store *ssa.Store + for len(pending) > 0 { + address := pending[len(pending)-1] + pending = pending[:len(pending)-1] + if seen[address] { + continue + } + seen[address] = true + users := address.Referrers() + if users == nil { + return nil + } + for _, user := range *users { + if user.Parent() == context.flow.function && !context.flow.blocks[user.Block()] { + continue + } + switch use := user.(type) { + case *ssa.DebugRef: + case *ssa.Store: + if use.Addr != address || address != allocation || use.Parent() != context.flow.function || store != nil { + return nil + } + store = use + case *ssa.UnOp: + if use.Op != token.MUL || use.X != address { + return nil + } + case *ssa.MakeClosure: + function, captures := ps6125SSACallee(use) + if function == nil { + return nil + } + for index, capture := range captures { + if capture == address { + pending = append(pending, function.FreeVars[index]) + } + } + default: + return nil + } + } + } + return store +} + +func ps6136AccessOwnerRoot(context *ps6125SSAContext, path ps6125AccessResult, owner *types.Named) ps6125SSAReference { + if !path.known { + return ps6125SSAReference{} + } + if root := ps6136OwnerRoot(context, path.access.root, owner); root.value != nil { + return root + } + for _, load := range path.access.loads { + if root := ps6136OwnerRoot(context, load, owner); root.value != nil { + return root + } + } + return ps6125SSAReference{} +} diff --git a/checks/ps6136_owner_root_test.go b/checks/ps6136_owner_root_test.go new file mode 100644 index 00000000..38006880 --- /dev/null +++ b/checks/ps6136_owner_root_test.go @@ -0,0 +1,58 @@ +package checks + +import ( + "go/types" + "testing" + + "golang.org/x/tools/go/ssa" +) + +func TestPS6136CapturedOwnerIdentity(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name, body string + known bool + }{ + {"readonly", `d:=input;mk:=func(){leaf(d)};mk()`, true}, + {"initialize_after_creation", `var d *owner;mk:=func(){leaf(d)};d=input;mk()`, true}, + {"initialize_after_call", `var d *owner;mk:=func(){leaf(d)};mk();d=input;_=d`, false}, + {"second_store", `d:=input;mk:=func(){leaf(d)};d=&owner{};mk()`, false}, + {"mixed_short_store", `d:=input;mk:=func(){leaf(d)};d,z:= &owner{},0;_=z;mk()`, false}, + {"captured_write", `d:=input;mk:=func(){d=&owner{};leaf(d)};mk()`, false}, + {"escaped_cell", `d:=input;mk:=func(){leaf(d)};expose(&d);mk()`, false}, + {"opaque_owner", `d:=opaque(input);mk:=func(){leaf(d)};mk()`, false}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + pkg := ps6125TestSSA(t, `package extent;type owner struct{width int};func leaf(*owner){};func expose(**owner){};func opaque(input *owner)*owner{return input};func root(input *owner){`+test.body+`}`) + function := pkg.Func("root") + root := ps6125NewSSAContext(function, nil, nil, 8) + var child *ps6125SSAContext + for _, call := range ps6125ContextCalls(root) { + if next := root.call(call); next != nil && len(next.flow.function.FreeVars) > 0 { + child = next + break + } + } + if child == nil { + t.Fatal("closed factory invocation unavailable") + } + var value ssa.Value + for _, call := range ps6125ContextCalls(child) { + if call.Call.StaticCallee() == pkg.Func("leaf") { + value = call.Call.Args[0] + break + } + } + if value == nil { + t.Fatal("owner leaf argument unavailable") + } + owner := pkg.Pkg.Scope().Lookup("owner").Type().(*types.Named) + got := ps6136OwnerRoot(child, value, owner) + want := root.reference(function.Params[0]) + if (got == want) != test.known { + t.Fatalf("captured owner identity=%v want %v", got == want, test.known) + } + }) + } +} diff --git a/checks/ps6136_owner_test.go b/checks/ps6136_owner_test.go new file mode 100644 index 00000000..99c4c90e --- /dev/null +++ b/checks/ps6136_owner_test.go @@ -0,0 +1,42 @@ +package checks + +import ( + "crypto/sha256" + "fmt" + "go/parser" + "go/token" + "os" + "testing" +) + +// These are complete byte-identical owner files, not hand-selected call shapes. +// This provenance test alone does not establish detector coverage or execute +// the owner's native backend. The analyzer replay must be tested separately. +func TestPS6136OwnerProvenance(t *testing.T) { + t.Parallel() + for _, fixture := range []struct{ path, digest string }{ + {"before/gpt.go.txt", "ad39b6e13d16557b749ecb9c1bbffe49eab6b0a90d33d3b8cd59f9d077bb35c7"}, + {"before/decoder.go.txt", "f13685b8af0cd797ea5b0c8d08983bcee90713f24b1b5fab5f2e8355ca7dbd17"}, + {"after/gpt.go.txt", "023ffa015e40f8405d9ddccff399a4c71e69dfda0818544f4c0d8691abaa4841"}, + {"after/decoder.go.txt", "b4a9355b118852e870a8b03444c99a7ef3f9bd8fb8b482efc0ebae35e88d1d16"}, + {"llamagpu.go.txt", "d404d5d843de7d9caabfe3e07cbbc40bca43723b1b55701217ae9a800bc468db"}, + {"cuda.go.txt", "8319c09049371d175d236fc2b6ca8d9336f88e9ae0db039de4afd3103d945d1c"}, + {"vulkan.go.txt", "610aa8af979ccdc779e05136cf01ee00f2eb6d99a1dde5b08c0239178cc419b5"}, + {"sampling_fastpath.go.txt", "e1f9e9d367ae6a1cca4d20be03f283622496b74e71def553eee7daba006c60e6"}, + } { + t.Run(fixture.path, func(t *testing.T) { + t.Parallel() + path := "testdata/ps6136-owner/" + fixture.path + data, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if got := fmt.Sprintf("%x", sha256.Sum256(data)); got != fixture.digest { + t.Fatalf("owner source changed: %s", got) + } + if _, err := parser.ParseFile(token.NewFileSet(), path, data, parser.AllErrors); err != nil { + t.Fatal(err) + } + }) + } +} diff --git a/checks/ps6136_owner_transfers.go b/checks/ps6136_owner_transfers.go new file mode 100644 index 00000000..4adb1349 --- /dev/null +++ b/checks/ps6136_owner_transfers.go @@ -0,0 +1,110 @@ +package checks + +import ( + "go/ast" + "go/types" + + "golang.org/x/tools/go/analysis" + "golang.org/x/tools/go/ssa" +) + +// ps6136OwnerTransfers inventories only direct source calls that forward an +// exact owner parameter/receiver, and fresh-constructor publication returns. +// Approvals cover just the owner operand, never a whole call/function. This +// closes no field effects or buffer extents by itself: the complete package's +// alias/effect/consumer gates remain mandatory. +func ps6136OwnerTransfers(pass *analysis.Pass, pkg *ssa.Package, owner *types.Named) map[ast.Node]bool { + if pass == nil || pkg == nil || owner == nil { + return nil + } + constructors := ps6136ConstructorInventory(pkg, owner) + if constructors == nil { + return nil + } + proved := make(map[ast.Node]bool) + isOwner := func(typ types.Type) bool { + return types.Identical(typ, owner) || types.Identical(typ, types.NewPointer(owner)) + } + // Never approve a compound operand: an outer source call/return must not + // hide an opaque escape or workspace read inside its argument expression. + simple := func(expression ast.Expr) bool { + _, identifier := ps2110Unparen(expression).(*ast.Ident) + return identifier + } + for _, file := range pass.Files { + ast.Inspect(file, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if !ok || call.Ellipsis.IsValid() { + return true + } + function, signature, ok := typedCallee(pass, call.Fun) + if !ok || signature.Variadic() || signature.TypeParams().Len() != 0 || signature.RecvTypeParams().Len() != 0 { + return true + } + callee := pkg.Prog.FuncValue(function) + if callee == nil || callee.Pkg != pkg || len(callee.Blocks) == 0 || callee.Syntax() == nil { + return true // No opaque owner-transfer semantics inferred. + } + arguments := call.Args + if signature.Recv() != nil && isOwner(signature.Recv().Type()) { + selector, ok := ps2110Unparen(call.Fun).(*ast.SelectorExpr) + if !ok { + return true + } + selection := pass.TypesInfo.Selections[selector] + if selection == nil { + return true + } + switch selection.Kind() { + case types.MethodVal: + if isOwner(pass.TypesInfo.TypeOf(selector.X)) && simple(selector.X) { + proved[selector.X] = true + } + case types.MethodExpr: + if len(arguments) == 0 || !isOwner(pass.TypesInfo.TypeOf(arguments[0])) || !simple(arguments[0]) { + return true + } + proved[arguments[0]] = true + arguments = arguments[1:] + default: + return true + } + } + if len(arguments) != signature.Params().Len() { + return true + } + for index, argument := range arguments { + if isOwner(signature.Params().At(index).Type()) && types.Identical(pass.TypesInfo.TypeOf(argument), signature.Params().At(index).Type()) && simple(argument) { + proved[argument] = true + } + } + return true + }) + } + for function := range constructors { + syntax := function.Syntax() + if syntax == nil { + return nil + } + var body *ast.BlockStmt + switch declaration := syntax.(type) { + case *ast.FuncDecl: + body = declaration.Body + case *ast.FuncLit: + body = declaration.Body + } + if body == nil { + return nil + } + ast.Inspect(body, func(node ast.Node) bool { + if _, nested := node.(*ast.FuncLit); nested { + return false + } + if returned, ok := node.(*ast.ReturnStmt); ok && len(returned.Results) > 0 && isOwner(pass.TypesInfo.TypeOf(returned.Results[0])) && simple(returned.Results[0]) { + proved[returned.Results[0]] = true + } + return true + }) + } + return proved +} diff --git a/checks/ps6136_owner_transfers_test.go b/checks/ps6136_owner_transfers_test.go new file mode 100644 index 00000000..38416324 --- /dev/null +++ b/checks/ps6136_owner_transfers_test.go @@ -0,0 +1,52 @@ +package checks + +import ( + "go/ast" + "go/parser" + "go/token" + "go/types" + "testing" + + "golang.org/x/tools/go/analysis" +) + +func TestPS6136ClosedSourceOwnerTransfers(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name, extra string + want bool + }{ + {"direct-source", "", true}, + {"alias", `func bad(d *Owner){other:=d;read(other)}`, false}, + {"opaque", `func opaque(*Owner);func bad(d *Owner){opaque(d)}`, false}, + {"interface-escape", `func sink(any){};func bad(d *Owner){sink(any(d))}`, false}, + {"compound-operand", `func opaque(*Owner);func wrap(d *Owner)*Owner{opaque(d);return &Owner{}};func bad(d *Owner){read(wrap(d))}`, false}, + {"captured-method", `func bad(d *Owner){fn:=d.Step;fn()}`, false}, + {"interior-pointer", `func bad(d *Owner){_= &((d.width))}`, false}, + {"whole-reset", `func bad(d *Owner){*d=Owner{}}`, false}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + fs := token.NewFileSet() + file, err := parser.ParseFile(fs, "transfers.go", `package transfers;type Owner struct{workspace []float32;width int};func fresh()*Owner{d:=&Owner{width:1};return d};func(d *Owner)Step(){read(d)};func read(d *Owner){_=d.width};`+test.extra, parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + info := &types.Info{Types: map[ast.Expr]types.TypeAndValue{}, Uses: map[*ast.Ident]types.Object{}, Defs: map[*ast.Ident]types.Object{}, Selections: map[*ast.SelectorExpr]*types.Selection{}} + pkg, err := (&types.Config{}).Check("transfers", fs, []*ast.File{file}, info) + if err != nil { + t.Fatal(err) + } + fixture := &analysisOwnerFixture{fs, []*ast.File{file}, info, pkg, nil} + ssaPkg := ps6136FixtureSSA(fixture) + owner := pkg.Scope().Lookup("Owner").Type().(*types.Named) + workspace := owner.Underlying().(*types.Struct).Field(0) + pass := &analysis.Pass{Fset: fs, Files: []*ast.File{file}, Pkg: pkg, TypesInfo: info} + transfers := ps6136OwnerTransfers(pass, ssaPkg, owner) + got := transfers != nil && ps6136ClosedObservations(pass, owner, workspace, transfers) + if got != test.want { + t.Fatalf("closed source transfers %v, want %v", got, test.want) + } + }) + } +} diff --git a/checks/ps6136_profile_recorder.go b/checks/ps6136_profile_recorder.go new file mode 100644 index 00000000..b9347e6c --- /dev/null +++ b/checks/ps6136_profile_recorder.go @@ -0,0 +1,257 @@ +package checks + +import ( + "go/ast" + "go/token" + "go/types" + + "golang.org/x/tools/go/analysis" +) + +// The profiler may add metadata, never a second native recorder choice. The +// operation method set must still select the exact base-adapter methods. +func ps6136ProfileRecorderFactory(pass *analysis.Pass, literal *ast.FuncLit, nativeID string, capacity, metadata types.Object, wrapper, base *types.Named, embedded, profile, nativeField *types.Var) bool { + if wrapper == nil || base == nil || embedded == nil || profile == nil || nativeField == nil || metadata == nil || capacity == nil { + return false + } + structure, ok := wrapper.Underlying().(*types.Struct) + if !ok || structure.NumFields() != 2 || structure.Field(0) != embedded || structure.Field(1) != profile || !embedded.Embedded() || !types.Identical(embedded.Type(), base) || !types.Identical(profile.Type(), types.NewPointer(metadata.Type())) { + return false + } + baseStructure, ok := base.Underlying().(*types.Struct) + if !ok || baseStructure.NumFields() != 1 || baseStructure.Field(0) != nativeField || !types.Identical(capacity.Type(), types.Typ[types.Int]) { + return false + } + return ps6136RecorderFactory(pass, literal, nativeID, nativeField.Type(), []types.Object{capacity}, func(expression ast.Expr, result types.Object) bool { + value, ok := expression.(*ast.CompositeLit) + if !ok || !types.Identical(pass.TypesInfo.TypeOf(value), wrapper) || len(value.Elts) != 2 { + return false + } + var nativeValue, metadataValue ast.Expr + for _, element := range value.Elts { + pair, ok := element.(*ast.KeyValueExpr) + if !ok { + return false + } + key, ok := pair.Key.(*ast.Ident) + if !ok { + return false + } + switch pass.TypesInfo.Uses[key] { + case embedded: + if nativeValue != nil { + return false + } + nativeValue = pair.Value + case profile: + if metadataValue != nil { + return false + } + metadataValue = pair.Value + default: + return false + } + } + baseValue, ok := nativeValue.(*ast.CompositeLit) + if !ok || !types.Identical(pass.TypesInfo.TypeOf(baseValue), base) || len(baseValue.Elts) != 1 { + return false + } + pair, ok := baseValue.Elts[0].(*ast.KeyValueExpr) + if !ok { + return false + } + key, ok := pair.Key.(*ast.Ident) + if !ok || pass.TypesInfo.Uses[key] != nativeField { + return false + } + identifier, ok := pair.Value.(*ast.Ident) + if !ok || pass.TypesInfo.Uses[identifier] != result { + return false + } + address, ok := metadataValue.(*ast.UnaryExpr) + if !ok || address.Op != token.AND { + return false + } + identifier, ok = address.X.(*ast.Ident) + return ok && pass.TypesInfo.Uses[identifier] == metadata + }) +} + +func ps6136SameAdapterMethods(profile, base *types.Named, methods []*types.Func) bool { + if profile == nil || base == nil || len(methods) == 0 { + return false + } + for _, method := range methods { + if method == nil { + return false + } + object, _, _ := types.LookupFieldOrMethod(profile, false, method.Pkg(), method.Name()) + if object != method { + return false + } + object, _, _ = types.LookupFieldOrMethod(base, false, method.Pkg(), method.Name()) + if object != method { + return false + } + } + return true +} + +// Exact source override/restore frame. It returns only the concrete assignment +// targets whose state changes are justified, never a function-wide exemption. +func ps6136ProfileOverrideFrame(pass *analysis.Pass, declaration *ast.FuncDecl, ops, primary, secondary, async *types.Var, stepID, dropID string, factoryProof func(*ast.FuncLit, types.Object, types.Object) bool) map[ast.Node]bool { + if pass == nil || declaration == nil || declaration.Body == nil || declaration.Recv == nil || len(declaration.Recv.List) != 1 || len(declaration.Recv.List[0].Names) != 1 || len(declaration.Body.List) != 14 || factoryProof == nil { + return nil + } + function, ok := pass.TypesInfo.Defs[declaration.Name].(*types.Func) + if !ok { + return nil + } + signature := function.Type().(*types.Signature) + if signature.Params().Len() != 3 || signature.Results().Len() != 3 { + return nil + } + receiver := pass.TypesInfo.Defs[declaration.Recv.List[0].Names[0]] + member := func(expression ast.Expr, field *types.Var) bool { + selector, ok := expression.(*ast.SelectorExpr) + if !ok { + return false + } + selection := pass.TypesInfo.Selections[selector] + base, ok := selector.X.(*ast.SelectorExpr) + if !ok || selection == nil || selection.Kind() != types.FieldVal || selection.Obj() != field { + return false + } + selection = pass.TypesInfo.Selections[base] + identifier, ok := base.X.(*ast.Ident) + return ok && pass.TypesInfo.Uses[identifier] == receiver && selection != nil && selection.Kind() == types.FieldVal && selection.Obj() == ops + } + identical := func(expression ast.Expr, object types.Object) bool { + identifier, ok := expression.(*ast.Ident) + return ok && pass.TypesInfo.Uses[identifier] == object + } + for _, statement := range declaration.Body.List[:4] { + guard, ok := statement.(*ast.IfStmt) + if !ok || guard.Init != nil || guard.Else != nil || len(guard.Body.List) != 1 { + return nil + } + returned, ok := guard.Body.List[0].(*ast.ReturnStmt) + if !ok || len(returned.Results) != 3 || !identical(returned.Results[0], types.Universe.Lookup("nil")) { + return nil + } + } + callStatement, ok := declaration.Body.List[4].(*ast.ExprStmt) + if !ok { + return nil + } + drop, ok := callStatement.X.(*ast.CallExpr) + if !ok || len(drop.Args) != 0 { + return nil + } + callee, _, typed := typedCallee(pass, drop.Fun) + selector, selected := drop.Fun.(*ast.SelectorExpr) + if !typed || ps6090FunctionID(callee) != dropID || !selected || !identical(selector.X, receiver) { + return nil + } + saved, ok := declaration.Body.List[5].(*ast.AssignStmt) + if !ok || saved.Tok != token.DEFINE || len(saved.Lhs) != 3 || len(saved.Rhs) != 3 { + return nil + } + fields := []*types.Var{primary, secondary, async} + previous := make([]types.Object, 3) + for index, field := range fields { + identifier, ok := saved.Lhs[index].(*ast.Ident) + if !ok || pass.TypesInfo.Defs[identifier] == nil || !member(saved.Rhs[index], field) { + return nil + } + previous[index] = pass.TypesInfo.Defs[identifier] + } + deferred, ok := declaration.Body.List[6].(*ast.DeferStmt) + if !ok || len(deferred.Call.Args) != 0 { + return nil + } + restore, ok := deferred.Call.Fun.(*ast.FuncLit) + if !ok || len(restore.Body.List) != 3 { + return nil + } + proved := make(map[ast.Node]bool, len(restore.Body.List)+3) + for index, statement := range restore.Body.List { + assignment, ok := statement.(*ast.AssignStmt) + if !ok || assignment.Tok != token.ASSIGN || len(assignment.Lhs) != 1 || len(assignment.Rhs) != 1 || !member(assignment.Lhs[0], fields[index]) || !identical(assignment.Rhs[0], previous[index]) { + return nil + } + proved[assignment.Lhs[0]] = true + } + metadataDeclaration, ok := declaration.Body.List[7].(*ast.DeclStmt) + if !ok { + return nil + } + general, ok := metadataDeclaration.Decl.(*ast.GenDecl) + if !ok || general.Tok != token.VAR || len(general.Specs) != 1 { + return nil + } + metadataValue, ok := general.Specs[0].(*ast.ValueSpec) + if !ok || len(metadataValue.Names) != 1 || len(metadataValue.Values) != 0 { + return nil + } + metadata := pass.TypesInfo.Defs[metadataValue.Names[0]] + factoryAssignment, ok := declaration.Body.List[8].(*ast.AssignStmt) + if !ok || factoryAssignment.Tok != token.ASSIGN || len(factoryAssignment.Lhs) != 1 || len(factoryAssignment.Rhs) != 1 || !member(factoryAssignment.Lhs[0], primary) { + return nil + } + literal, ok := factoryAssignment.Rhs[0].(*ast.FuncLit) + if !ok || !factoryProof(literal, signature.Params().At(2), metadata) { + return nil + } + proved[factoryAssignment.Lhs[0]] = true + alias, ok := declaration.Body.List[9].(*ast.AssignStmt) + if !ok || alias.Tok != token.ASSIGN || len(alias.Lhs) != 1 || len(alias.Rhs) != 1 || !member(alias.Lhs[0], secondary) || !member(alias.Rhs[0], primary) { + return nil + } + proved[alias.Lhs[0]] = true + synchronous, ok := declaration.Body.List[10].(*ast.AssignStmt) + if !ok || synchronous.Tok != token.ASSIGN || len(synchronous.Lhs) != 1 || len(synchronous.Rhs) != 1 || !member(synchronous.Lhs[0], async) || !identical(synchronous.Rhs[0], types.Universe.Lookup("false")) { + return nil + } + proved[synchronous.Lhs[0]] = true + step, ok := declaration.Body.List[11].(*ast.AssignStmt) + if !ok || step.Tok != token.DEFINE || len(step.Lhs) != 2 || len(step.Rhs) != 1 { + return nil + } + call, ok := step.Rhs[0].(*ast.CallExpr) + if !ok || len(call.Args) != 2 || !identical(call.Args[0], signature.Params().At(0)) || !identical(call.Args[1], signature.Params().At(1)) { + return nil + } + callee, _, typed = typedCallee(pass, call.Fun) + selector, selected = call.Fun.(*ast.SelectorExpr) + if !typed || ps6090FunctionID(callee) != stepID || !selected || !identical(selector.X, receiver) { + return nil + } + // No late operations may write the saved callback cells or bypass restoration. + guard, ok := declaration.Body.List[12].(*ast.IfStmt) + if !ok || guard.Init != nil || guard.Else != nil || len(guard.Body.List) != 1 { + return nil + } + logits, logitsOK := step.Lhs[0].(*ast.Ident) + errorName, errorOK := step.Lhs[1].(*ast.Ident) + if !logitsOK || !errorOK || pass.TypesInfo.Defs[logits] == nil || pass.TypesInfo.Defs[errorName] == nil { + return nil + } + comparison, ok := guard.Cond.(*ast.BinaryExpr) + if !ok || comparison.Op != token.NEQ || !identical(comparison.X, pass.TypesInfo.Defs[errorName]) || !identical(comparison.Y, types.Universe.Lookup("nil")) { + return nil + } + failure, ok := guard.Body.List[0].(*ast.ReturnStmt) + if !ok || len(failure.Results) != 3 || !identical(failure.Results[0], types.Universe.Lookup("nil")) || !identical(failure.Results[2], pass.TypesInfo.Defs[errorName]) { + return nil + } + emptyProfile, ok := failure.Results[1].(*ast.CompositeLit) + if !ok || len(emptyProfile.Elts) != 0 || !types.Identical(pass.TypesInfo.TypeOf(emptyProfile), metadata.Type()) { + return nil + } + publication, ok := declaration.Body.List[13].(*ast.ReturnStmt) + if !ok || len(publication.Results) != 3 || !identical(publication.Results[0], pass.TypesInfo.Defs[logits]) || !identical(publication.Results[1], metadata) || !identical(publication.Results[2], types.Universe.Lookup("nil")) { + return nil + } + return proved +} diff --git a/checks/ps6136_profile_recorder_test.go b/checks/ps6136_profile_recorder_test.go new file mode 100644 index 00000000..1959121b --- /dev/null +++ b/checks/ps6136_profile_recorder_test.go @@ -0,0 +1,102 @@ +package checks + +import ( + "go/ast" + "go/token" + "go/types" + "testing" + + "golang.org/x/tools/go/analysis" +) + +func TestPS6136AuthenticProfileRecorderBinding(t *testing.T) { + t.Parallel() + for _, mutation := range []string{"original", "wrong_factory", "reversed_error", "wrong_native_field", "wrong_metadata", "late_factory", "wrong_restore", "late_frame"} { + t.Run(mutation, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, "before", true, true) + if err != nil { + t.Fatal(err) + } + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg} + var declaration *ast.FuncDecl + for _, file := range fixture.files { + for _, decl := range file.Decls { + if function, ok := decl.(*ast.FuncDecl); ok && function.Name.Name == "ProfileMetalStep" { + declaration = function + } + } + } + if declaration == nil { + t.Fatal("authentic profile method missing") + } + function := fixture.info.Defs[declaration.Name].(*types.Func) + capacity := function.Type().(*types.Signature).Params().At(2) + var metadata types.Object + var literal *ast.FuncLit + ast.Inspect(declaration.Body, func(node ast.Node) bool { + if value, ok := node.(*ast.ValueSpec); ok && len(value.Names) == 1 && value.Names[0].Name == "profile" { + metadata = fixture.info.Defs[value.Names[0]] + } + if assignment, ok := node.(*ast.AssignStmt); ok && len(assignment.Rhs) == 1 { + if candidate, ok := assignment.Rhs[0].(*ast.FuncLit); ok { + literal = candidate + } + } + return true + }) + if literal == nil || metadata == nil { + t.Fatal("actual profile callback flow missing") + } + wrapper := fixture.pkg.Scope().Lookup("mProfileRec").Type().(*types.Named) + base := fixture.pkg.Scope().Lookup("mRec").Type().(*types.Named) + structure := wrapper.Underlying().(*types.Struct) + field := base.Underlying().(*types.Struct).Field(0) + id := "github.com/jxsl13/goai/backend/metal.NewProfilingRecorder" + switch mutation { + case "wrong_factory": + id = "github.com/jxsl13/goai/backend/metal.NewRecorder" + case "reversed_error": + literal.Body.List[1].(*ast.IfStmt).Cond.(*ast.BinaryExpr).Op = token.EQL + case "wrong_native_field": + field = types.NewVar(token.NoPos, fixture.pkg, "other", field.Type()) + case "wrong_metadata": + metadata = types.NewVar(token.NoPos, fixture.pkg, "other", metadata.Type()) + case "late_factory": + literal.Body.List = append(literal.Body.List[:2], &ast.ExprStmt{X: literal.Body.List[0].(*ast.AssignStmt).Rhs[0]}, literal.Body.List[2]) + case "wrong_restore": + restore := declaration.Body.List[6].(*ast.DeferStmt).Call.Fun.(*ast.FuncLit) + restore.Body.List[0].(*ast.AssignStmt).Rhs[0] = restore.Body.List[1].(*ast.AssignStmt).Rhs[0] + case "late_frame": + declaration.Body.List = append(declaration.Body.List[:13], &ast.ExprStmt{X: literal.Body.List[0].(*ast.AssignStmt).Rhs[0]}, declaration.Body.List[13]) + } + factoryValid := mutation == "original" || mutation == "wrong_restore" || mutation == "late_frame" + if got := ps6136ProfileRecorderFactory(pass, literal, id, capacity, metadata, wrapper, base, structure.Field(0), structure.Field(1), field); got != factoryValid { + t.Fatalf("native/base/metadata identity proof %v", got) + } + owner := fixture.pkg.Scope().Lookup("Decoder").Type() + lookup := func(typ types.Type, name string) *types.Var { + object, _, _ := types.LookupFieldOrMethod(typ, true, fixture.pkg, name) + return object.(*types.Var) + } + ops := lookup(owner, "ops") + proof := ps6136ProfileOverrideFrame(pass, declaration, ops, lookup(ops.Type(), "newRecorder"), lookup(ops.Type(), "newDecodeRecorder"), lookup(ops.Type(), "asyncEncode"), "github.com/jxsl13/goai/llamagpu.Decoder.Step", "github.com/jxsl13/goai/llamagpu.Decoder.dropPending", func(callback *ast.FuncLit, actualCapacity, actualMetadata types.Object) bool { + if mutation == "wrong_metadata" { + actualMetadata = metadata + } + return ps6136ProfileRecorderFactory(pass, callback, id, actualCapacity, actualMetadata, wrapper, base, structure.Field(0), structure.Field(1), field) + }) + if (proof != nil) != (mutation == "original") { + t.Fatalf("exact deferred override/restore frame proof %v", proof != nil) + } + var methods []*types.Func + for _, name := range []string{"MatMul", "AddBias"} { + object, _, _ := types.LookupFieldOrMethod(base, false, fixture.pkg, name) + methods = append(methods, object.(*types.Func)) + } + if !ps6136SameAdapterMethods(wrapper, base, methods) { + t.Fatal("profiling wrapper does not inherit exactly the proved native adapters") + } + }) + } +} diff --git a/checks/ps6136_recorder_binding.go b/checks/ps6136_recorder_binding.go new file mode 100644 index 00000000..8b9464c9 --- /dev/null +++ b/checks/ps6136_recorder_binding.go @@ -0,0 +1,124 @@ +package checks + +import ( + "go/ast" + "go/token" + "go/types" + + "golang.org/x/tools/go/analysis" +) + +func (selection *ps6136Selection) recorderBinding(pass *analysis.Pass, recorderField *types.Var, nativeID string, wrapper *types.Named, field *types.Var) bool { + if selection == nil || recorderField == nil || selection.ops == nil { + return false + } + snapshot := ps6136InitialOwnerField(selection.owner, selection.ops) + if snapshot.value == nil { + return false + } + structure, ok := selection.ops.Type().Underlying().(*types.Struct) + if !ok { + return false + } + index := -1 + for candidate := 0; candidate < structure.NumFields(); candidate++ { + if structure.Field(candidate) == recorderField { + index = candidate + } + } + if index < 0 { + return false + } + origin := snapshot.context.structField(snapshot.value, index) + if origin.value == nil || origin.context == nil { + return false + } + function, captures := ps6125SSACallee(origin.value) + if function == nil || len(captures) != 0 || len(function.FreeVars) != 0 { + return false + } + literal, ok := function.Syntax().(*ast.FuncLit) + return ok && ps6136RecorderLiteral(pass, literal, nativeID, wrapper, field) +} + +// Bind the selected backend's source literal to its exact native recorder. +// This closes the factory's argument/result flow, not the native API's access, +// scheduling, synchronization or lifetime semantics. +func ps6136RecorderLiteral(pass *analysis.Pass, literal *ast.FuncLit, nativeID string, wrapper *types.Named, field *types.Var) bool { + if wrapper == nil || field == nil { + return false + } + return ps6136RecorderFactory(pass, literal, nativeID, field.Type(), nil, func(expression ast.Expr, result types.Object) bool { + value, ok := expression.(*ast.CompositeLit) + if !ok || !types.Identical(pass.TypesInfo.TypeOf(value), wrapper) || len(value.Elts) != 1 { + return false + } + structure, ok := wrapper.Underlying().(*types.Struct) + if !ok || structure.NumFields() != 1 || structure.Field(0) != field { + return false + } + identifier, ok := value.Elts[0].(*ast.Ident) + return ok && pass.TypesInfo.Uses[identifier] == result + }) +} + +func ps6136RecorderFactory(pass *analysis.Pass, literal *ast.FuncLit, nativeID string, nativeType types.Type, arguments []types.Object, resultProof func(ast.Expr, types.Object) bool) bool { + if pass == nil || literal == nil || nativeType == nil || resultProof == nil || len(literal.Body.List) != 3 { + return false + } + signature, ok := pass.TypesInfo.TypeOf(literal).(*types.Signature) + if !ok || signature.Variadic() || signature.Params().Len() != 0 || signature.Results().Len() != 2 || !types.Identical(signature.Results().At(1).Type(), types.Universe.Lookup("error").Type()) { + return false + } + assignment, ok := literal.Body.List[0].(*ast.AssignStmt) + if !ok || assignment.Tok != token.DEFINE || len(assignment.Lhs) != 2 || len(assignment.Rhs) != 1 { + return false + } + result, ok := assignment.Lhs[0].(*ast.Ident) + if !ok || pass.TypesInfo.Defs[result] == nil { + return false + } + errorName, ok := assignment.Lhs[1].(*ast.Ident) + if !ok || pass.TypesInfo.Defs[errorName] == nil { + return false + } + call, ok := assignment.Rhs[0].(*ast.CallExpr) + if !ok || len(call.Args) != len(arguments) || call.Ellipsis.IsValid() { + return false + } + callee, nativeSignature, ok := typedCallee(pass, call.Fun) + if !ok || ps6090FunctionID(callee) != nativeID || nativeSignature.Recv() != nil || nativeSignature.Variadic() || nativeSignature.TypeParams().Len() != 0 || nativeSignature.Params().Len() != len(arguments) || nativeSignature.Results().Len() != 2 || !types.Identical(nativeSignature.Results().At(0).Type(), nativeType) || !types.Identical(nativeSignature.Results().At(1).Type(), types.Universe.Lookup("error").Type()) { + return false + } + for index, object := range arguments { + identifier, ok := call.Args[index].(*ast.Ident) + if !ok || object == nil || pass.TypesInfo.Uses[identifier] != object || !types.Identical(object.Type(), nativeSignature.Params().At(index).Type()) { + return false + } + } + guard, ok := literal.Body.List[1].(*ast.IfStmt) + if !ok || guard.Init != nil || guard.Else != nil || len(guard.Body.List) != 1 { + return false + } + condition, ok := guard.Cond.(*ast.BinaryExpr) + if !ok || condition.Op != token.NEQ { + return false + } + identical := func(expression ast.Expr, object types.Object) bool { + identifier, ok := expression.(*ast.Ident) + return ok && pass.TypesInfo.Uses[identifier] == object + } + nilValue := func(expression ast.Expr) bool { return identical(expression, types.Universe.Lookup("nil")) } + if !identical(condition.X, pass.TypesInfo.Defs[errorName]) || !nilValue(condition.Y) { + return false + } + failure, ok := guard.Body.List[0].(*ast.ReturnStmt) + if !ok || len(failure.Results) != 2 || !nilValue(failure.Results[0]) || !identical(failure.Results[1], pass.TypesInfo.Defs[errorName]) { + return false + } + publication, ok := literal.Body.List[2].(*ast.ReturnStmt) + if !ok || len(publication.Results) != 2 || !nilValue(publication.Results[1]) { + return false + } + return resultProof(publication.Results[0], pass.TypesInfo.Defs[result]) +} diff --git a/checks/ps6136_recorder_binding_test.go b/checks/ps6136_recorder_binding_test.go new file mode 100644 index 00000000..c964eaa4 --- /dev/null +++ b/checks/ps6136_recorder_binding_test.go @@ -0,0 +1,67 @@ +package checks + +import ( + "go/ast" + "go/token" + "go/types" + "testing" + + "golang.org/x/tools/go/analysis" +) + +func TestPS6136AuthenticRecorderFactoryBoundaries(t *testing.T) { + t.Parallel() + for _, entry := range []string{"New", "NewGPT"} { + for _, mutation := range []string{"original", "wrong_native", "reversed_error", "late_factory", "discard_native"} { + t.Run(entry+"/"+mutation, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, "before", true, true) + if err != nil { + t.Fatal(err) + } + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg} + wrapper := fixture.pkg.Scope().Lookup("mRec").Type().(*types.Named) + field := wrapper.Underlying().(*types.Struct).Field(0) + var literal *ast.FuncLit + for _, file := range fixture.files { + for _, declaration := range file.Decls { + function, ok := declaration.(*ast.FuncDecl) + if !ok || function.Name.Name != entry { + continue + } + ast.Inspect(function.Body, func(node ast.Node) bool { + pair, ok := node.(*ast.KeyValueExpr) + if !ok { + return true + } + key, ok := pair.Key.(*ast.Ident) + if ok && key.Name == "newRecorder" { + literal, _ = pair.Value.(*ast.FuncLit) + } + return true + }) + } + } + if literal == nil { + t.Fatal("pinned selected constructor factory missing") + } + id := "github.com/jxsl13/goai/backend/metal.NewRecorder" + switch mutation { + case "wrong_native": + id = "github.com/jxsl13/goai/backend/metal.NewConcurrentRecorder" + case "reversed_error": + literal.Body.List[1].(*ast.IfStmt).Cond.(*ast.BinaryExpr).Op = token.EQL + case "late_factory": + literal.Body.List = append(literal.Body.List[:2], &ast.ExprStmt{X: literal.Body.List[0].(*ast.AssignStmt).Rhs[0]}, literal.Body.List[2]) + case "discard_native": + nilValue := &ast.Ident{Name: "nil"} + fixture.info.Uses[nilValue] = types.Universe.Lookup("nil") + literal.Body.List[2].(*ast.ReturnStmt).Results[0].(*ast.CompositeLit).Elts[0] = nilValue + } + if got := ps6136RecorderLiteral(pass, literal, id, wrapper, field); got != (mutation == "original") { + t.Fatalf("exact factory binding %v", got) + } + }) + } + } +} diff --git a/checks/ps6136_release.go b/checks/ps6136_release.go new file mode 100644 index 00000000..f52d344a --- /dev/null +++ b/checks/ps6136_release.go @@ -0,0 +1,94 @@ +package checks + +import ( + "go/ast" + "go/token" + "go/types" + + "golang.org/x/tools/go/analysis" +) + +// ps6136ReleaseList proves the source loop releases every non-nil retained +// element once and immediately drops the list. This is not a universal +// lifecycle proof: the caller must close all list observations, aliases and +// other release paths, and bind the reviewed native release implementation. +func ps6136ReleaseList(pass *analysis.Pass, declaration *ast.FuncDecl, receiver types.Object, list *types.Var, release *types.Func) map[ast.Node]bool { + if pass == nil || declaration == nil || declaration.Body == nil || receiver == nil || list == nil || release == nil { + return nil + } + var approved map[ast.Node]bool + for index, statement := range declaration.Body.List { + loop, ok := statement.(*ast.RangeStmt) + if !ok || loop.Tok != token.DEFINE || index+1 >= len(declaration.Body.List) || !ps6136DirectField(pass, loop.X, receiver, list) { + continue + } + if approved != nil || len(loop.Body.List) != 1 { + return nil + } + // The cleanup must be reached on every normally returning invocation. + // Permit straight-line expression prefixes (the authentic shared owner + // drains pending work here), but no conditional return, label/goto, + // nested control flow, deferred or asynchronous prefix execution. + for _, prefix := range declaration.Body.List[:index] { + if _, ok := prefix.(*ast.ExprStmt); !ok { + return nil + } + } + value, ok := loop.Value.(*ast.Ident) + if !ok || pass.TypesInfo.Defs[value] == nil { + return nil + } + guard, ok := loop.Body.List[0].(*ast.IfStmt) + if !ok || guard.Init != nil || guard.Else != nil || len(guard.Body.List) != 1 { + return nil + } + condition, ok := guard.Cond.(*ast.BinaryExpr) + if !ok || condition.Op != token.NEQ || !ps6136ObjectExpr(pass, condition.X, pass.TypesInfo.Defs[value]) || !ps6136NilExpr(pass, condition.Y) { + return nil + } + expression, ok := guard.Body.List[0].(*ast.ExprStmt) + if !ok { + return nil + } + call, ok := expression.X.(*ast.CallExpr) + if !ok || len(call.Args) != 0 || call.Ellipsis.IsValid() { + return nil + } + method, ok := call.Fun.(*ast.SelectorExpr) + if !ok || !ps6136ObjectExpr(pass, method.X, pass.TypesInfo.Defs[value]) { + return nil + } + selection := pass.TypesInfo.Selections[method] + if selection == nil || selection.Obj() != release { + return nil + } + signature, ok := release.Type().(*types.Signature) + if !ok || signature.Variadic() || signature.Params().Len() != 0 || signature.Results().Len() != 0 { + return nil + } + clear, ok := declaration.Body.List[index+1].(*ast.AssignStmt) + if !ok || clear.Tok != token.ASSIGN || len(clear.Lhs) != 1 || len(clear.Rhs) != 1 || !ps6136DirectField(pass, clear.Lhs[0], receiver, list) || !ps6136NilExpr(pass, clear.Rhs[0]) { + return nil + } + approved = map[ast.Node]bool{loop.X: true, clear.Lhs[0]: true} + } + return approved +} + +func ps6136ObjectExpr(pass *analysis.Pass, expression ast.Expr, object types.Object) bool { + identifier, ok := expression.(*ast.Ident) + return ok && object != nil && pass.TypesInfo.Uses[identifier] == object +} + +func ps6136NilExpr(pass *analysis.Pass, expression ast.Expr) bool { + return ps6136ObjectExpr(pass, expression, types.Universe.Lookup("nil")) +} + +func ps6136DirectField(pass *analysis.Pass, expression ast.Expr, receiver types.Object, field *types.Var) bool { + selector, ok := expression.(*ast.SelectorExpr) + if !ok || !ps6136ObjectExpr(pass, selector.X, receiver) { + return false + } + selection := pass.TypesInfo.Selections[selector] + return selection != nil && selection.Kind() == types.FieldVal && selection.Obj() == field +} diff --git a/checks/ps6136_release_test.go b/checks/ps6136_release_test.go new file mode 100644 index 00000000..ab624780 --- /dev/null +++ b/checks/ps6136_release_test.go @@ -0,0 +1,96 @@ +package checks + +import ( + "go/ast" + "go/parser" + "go/token" + "go/types" + "testing" + + "golang.org/x/tools/go/analysis" +) + +func TestPS6136AuthenticReleaseLists(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + t.Run(revision, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, revision, false, true) + if err != nil { + t.Fatal(err) + } + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg} + for _, ownerName := range []string{"GPTDecoder", "Decoder"} { + owner := fixture.pkg.Scope().Lookup(ownerName).Type().(*types.Named) + list, _, _ := types.LookupFieldOrMethod(owner, true, fixture.pkg, "all") + element := list.Type().Underlying().(*types.Slice).Elem() + release, _, _ := types.LookupFieldOrMethod(element, true, fixture.pkg, "Release") + found := false + for _, file := range fixture.files { + for _, declaration := range file.Decls { + function, ok := declaration.(*ast.FuncDecl) + if !ok || function.Recv == nil || function.Name.Name != "Release" || !types.Identical(fixture.info.TypeOf(function.Recv.List[0].Type), types.NewPointer(owner)) { + continue + } + found = true + receiver := fixture.info.Defs[function.Recv.List[0].Names[0]] + if ps6136ReleaseList(pass, function, receiver, list.(*types.Var), release.(*types.Func)) == nil { + t.Fatalf("authentic %s retained-list release unproved", ownerName) + } + } + } + if !found { + t.Fatalf("authentic %s release absent", ownerName) + } + } + }) + } +} + +func TestPS6136ReleaseListBoundaries(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name, body string + valid bool + }{ + {"canonical", `for _,b:=range d.all{if b!=nil{b.Release()}};d.all=nil`, true}, + {"early_return", `if flag{return};for _,b:=range d.all{if b!=nil{b.Release()}};d.all=nil`, false}, + {"goto_bypass", `goto done;for _,b:=range d.all{if b!=nil{b.Release()}};d.all=nil;done:return`, false}, + {"conditional_clear", `for _,b:=range d.all{if b!=nil{b.Release()}};if flag{d.all=nil}`, false}, + {"break", `for _,b:=range d.all{if b!=nil{b.Release();break}};d.all=nil`, false}, + {"double_release", `for _,b:=range d.all{if b!=nil{b.Release();b.Release()}};d.all=nil`, false}, + {"wrong_release", `for _,b:=range d.all{if b!=nil{b.Close()}};d.all=nil`, false}, + {"filtered", `for i,b:=range d.all{if i==0&&b!=nil{b.Release()}};d.all=nil`, false}, + {"wrong_list", `for _,b:=range d.other{if b!=nil{b.Release()}};d.all=nil`, false}, + {"no_clear", `for _,b:=range d.all{if b!=nil{b.Release()}}`, false}, + {"clear_wrong_list", `for _,b:=range d.all{if b!=nil{b.Release()}};d.other=nil`, false}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + fs := token.NewFileSet() + file, err := parser.ParseFile(fs, "release.go", `package ownership;var flag bool;type buffer interface{Release();Close()};type owner struct{all,other []buffer};func(d *owner)release(){`+test.body+`}`, parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + info := &types.Info{Types: map[ast.Expr]types.TypeAndValue{}, Uses: map[*ast.Ident]types.Object{}, Defs: map[*ast.Ident]types.Object{}, Selections: map[*ast.SelectorExpr]*types.Selection{}} + pkg, err := (&types.Config{}).Check("ownership", fs, []*ast.File{file}, info) + if err != nil { + t.Fatal(err) + } + owner := pkg.Scope().Lookup("owner").Type().(*types.Named) + list, _, _ := types.LookupFieldOrMethod(owner, true, pkg, "all") + release, _, _ := types.LookupFieldOrMethod(pkg.Scope().Lookup("buffer").Type(), true, pkg, "Release") + var declaration *ast.FuncDecl + for _, node := range file.Decls { + if function, ok := node.(*ast.FuncDecl); ok { + declaration = function + } + } + pass := &analysis.Pass{Fset: fs, Files: []*ast.File{file}, TypesInfo: info, Pkg: pkg} + approved := ps6136ReleaseList(pass, declaration, info.Defs[declaration.Recv.List[0].Names[0]], list.(*types.Var), release.(*types.Func)) + if (approved != nil) != test.valid { + t.Fatalf("closed retained-list release=%v want %v", approved != nil, test.valid) + } + }) + } +} diff --git a/checks/ps6136_residency.go b/checks/ps6136_residency.go new file mode 100644 index 00000000..c8d879f1 --- /dev/null +++ b/checks/ps6136_residency.go @@ -0,0 +1,45 @@ +package checks + +// ps6136Residency describes geometry only. Callers must independently prove +// ownership, complete logical-use coverage and the reviewed workload policy. +// Numeric profile values are not inferred runtime dimensions. +type ps6136Residency struct { + retained ps6125Extent + common ps6125Extent + rows *ps6125ExtentIdentity + width *ps6125ExtentIdentity +} + +func ps6136OutputResidency(rows, width *ps6125ExtentIdentity) (ps6136Residency, bool) { + if rows == nil || width == nil || rows == width || rows.pool != width.pool { + return ps6136Residency{}, false + } + root := func(identity *ps6125ExtentIdentity) *ps6125ExtentIdentity { + for identity.parent != nil { + identity = identity.parent + } + return identity + } + if root(rows) != root(width) { + return ps6136Residency{}, false + } + common := ps6125MultiplyExtents(ps6125ConstantExtent(4), ps6125SymbolicExtent(width)) + retained := ps6125MultiplyExtents(common, ps6125SymbolicExtent(rows)) + return ps6136Residency{retained: retained, common: common, rows: rows, width: width}, retained.known +} + +func (r ps6136Residency) evaluate(rows, width, limit int64) (retained, common, idle, amplification int64, known bool) { + if rows <= 1 || width <= 0 || r.rows == nil || r.width == nil { + return 0, 0, 0, 0, false + } + dimensions := map[*ps6125ExtentIdentity]int64{r.rows: rows, r.width: width} + retained, valid := ps6125EvaluateExtent(r.retained, dimensions, limit) + if !valid { + return 0, 0, 0, 0, false + } + common, valid = ps6125EvaluateExtent(r.common, dimensions, limit) + if !valid || retained <= common { + return 0, 0, 0, 0, false + } + return retained, common, retained - common, rows, true +} diff --git a/checks/ps6136_residency_test.go b/checks/ps6136_residency_test.go new file mode 100644 index 00000000..ee03785c --- /dev/null +++ b/checks/ps6136_residency_test.go @@ -0,0 +1,37 @@ +package checks + +import ( + "go/token" + "go/types" + "math" + "testing" +) + +func TestPS6136Residency(t *testing.T) { + t.Parallel() + var pool ps6125ExtentPool + owner := types.NewVar(token.NoPos, nil, "decoder", types.Typ[types.Int]) + rows := pool.identity(owner, []*types.Var{types.NewField(token.NoPos, nil, "context", types.Typ[types.Int], false)}, false) + width := pool.identity(owner, []*types.Var{types.NewField(token.NoPos, nil, "vocabulary", types.Typ[types.Int], false)}, false) + r, valid := ps6136OutputResidency(rows, width) + if !valid { + t.Fatal("same-owner distinct geometry rejected") + } + retained, common, idle, amplification, known := r.evaluate(1024, 50257, math.MaxInt64) + if !known || retained != 205852672 || common != 201028 || idle != 205651644 || amplification != 1024 { + t.Fatalf("owner profile: %d %d %d %d %v", retained, common, idle, amplification, known) + } + for _, profile := range [][3]int64{{0, 50257, math.MaxInt64}, {1, 50257, math.MaxInt64}, {1024, 0, math.MaxInt64}, {-1, 50257, math.MaxInt64}, {1024, -1, math.MaxInt64}, {1024, 50257, 1000}, {math.MaxInt64, 50257, math.MaxInt64}, {1024, math.MaxInt64, math.MaxInt64}} { + if _, _, _, _, known := r.evaluate(profile[0], profile[1], profile[2]); known { + t.Fatalf("invalid/overflow profile accepted: %v", profile) + } + } + var foreign ps6125ExtentPool + other := foreign.identity(owner, nil, false) + otherOwner := pool.identity(types.NewVar(token.NoPos, nil, "otherDecoder", types.Typ[types.Int]), nil, false) + for _, pair := range [][2]*ps6125ExtentIdentity{{nil, width}, {rows, nil}, {rows, rows}, {rows, other}, {rows, otherOwner}} { + if _, valid := ps6136OutputResidency(pair[0], pair[1]); valid { + t.Fatal("unknown/contradictory geometry accepted") + } + } +} diff --git a/checks/ps6136_retention.go b/checks/ps6136_retention.go new file mode 100644 index 00000000..d6145a6d --- /dev/null +++ b/checks/ps6136_retention.go @@ -0,0 +1,184 @@ +package checks + +import ( + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136FactoryRetention proves the successful result is appended exactly once +// to this owner's retention list and stored only in its returned fresh slot. +// It does not infer native exclusivity, successful execution or final release. +func ps6136FactoryRetention(context *ps6125SSAContext, backend *ssa.Call, owner ps6125SSAReference, ownerType *types.Named, list, slot *types.Var) bool { + if context == nil || backend == nil || owner.value == nil || ownerType == nil || list == nil || slot == nil { + return false + } + paths := ps6125AccessPaths{flow: context.flow} + var result *ssa.Extract + users := backend.Referrers() + if users == nil { + return false + } + for _, user := range *users { + if extract, ok := user.(*ssa.Extract); ok && extract.Index == 0 { + if result != nil { + return false + } + result = extract + } + } + if result == nil { + return false + } + var appendCall *ssa.Call + var array *ssa.Alloc + var returnedSlot *ssa.Alloc + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + if returned, ok := instruction.(*ssa.Return); ok { + allocation, fields, known := context.returnedFields(returned, 0) + if !known { + return false + } + if fields[slot].value == result { + candidate, ok := allocation.value.(*ssa.Alloc) + if !ok || allocation.context != context || returnedSlot != nil { + return false + } + returnedSlot = candidate + } + } + store, ok := instruction.(*ssa.Store) + if !ok { + continue + } + path := paths.resolve(store.Addr) + if !path.known || len(path.access.fields) != 1 || path.access.fields[0] != list { + continue + } + if ps6136AccessOwnerRoot(context, path, ownerType) != owner || appendCall != nil { + return false + } + call, ok := store.Val.(*ssa.Call) + if !ok { + return false + } + builtin, ok := call.Call.Value.(*ssa.Builtin) + if !ok || builtin.Name() != "append" || len(call.Call.Args) != 2 { + return false + } + current := paths.resolve(call.Call.Args[0]) + if !current.known || len(current.access.fields) != 1 || current.access.fields[0] != list || ps6136AccessOwnerRoot(context, current, ownerType) != owner { + return false + } + candidate := ps6136SingleAppend(context, call.Call.Args[1], call, result) + if candidate == nil { + return false + } + appendCall, array = call, candidate + } + } + if appendCall == nil || returnedSlot == nil { + return false + } + resultUsers := result.Referrers() + if resultUsers == nil { + return false + } + for _, user := range *resultUsers { + if user.Parent() == context.flow.function && !context.flow.blocks[user.Block()] { + continue + } + switch use := user.(type) { + case *ssa.DebugRef: + case *ssa.Store: + switch address := use.Addr.(type) { + case *ssa.IndexAddr: + if address.X != array { + return false + } + case *ssa.FieldAddr: + path := paths.resolve(address) + if address.X != returnedSlot || !path.known || len(path.access.fields) != 1 || path.access.fields[0] != slot { + return false + } + default: + return false + } + default: + return false // return aliases, extra helpers, captures, stores and conversions + } + } + return true +} + +func ps6136SingleAppend(context *ps6125SSAContext, value ssa.Value, call *ssa.Call, result ssa.Value) *ssa.Alloc { + slice, ok := value.(*ssa.Slice) + if !ok || slice.Low != nil || slice.High != nil || slice.Max != nil { + return nil + } + array, ok := slice.X.(*ssa.Alloc) + if !ok || array.Parent() != context.flow.function { + return nil + } + pointer, ok := array.Type().(*types.Pointer) + if !ok { + return nil + } + shape, ok := pointer.Elem().Underlying().(*types.Array) + if !ok || shape.Len() != 1 || !types.Identical(shape.Elem(), result.Type()) { + return nil + } + users := array.Referrers() + if users == nil { + return nil + } + stores, slices := 0, 0 + for _, user := range *users { + switch use := user.(type) { + case *ssa.DebugRef: + case *ssa.Slice: + if use != slice { + return nil + } + slices++ + case *ssa.IndexAddr: + index := context.scalar(use.Index) + if index.state != ps6125Integer || !ps6125SameExtent(index.extent, ps6125ConstantExtent(0)) { + return nil + } + addressUsers := use.Referrers() + if addressUsers == nil { + return nil + } + for _, addressUser := range *addressUsers { + if _, debug := addressUser.(*ssa.DebugRef); debug { + continue + } + store, ok := addressUser.(*ssa.Store) + if !ok || store.Addr != use || store.Val != result || !context.flow.instructionDominates(store, call) { + return nil + } + stores++ + } + default: + return nil + } + } + if stores != 1 || slices != 1 { + return nil + } + sliceUsers := slice.Referrers() + if sliceUsers == nil { + return nil + } + for _, user := range *sliceUsers { + if _, debug := user.(*ssa.DebugRef); !debug && user != call { + return nil + } + } + return array +} diff --git a/checks/ps6136_returned_callable.go b/checks/ps6136_returned_callable.go new file mode 100644 index 00000000..7282685b --- /dev/null +++ b/checks/ps6136_returned_callable.go @@ -0,0 +1,120 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136Call adds only closed, source-visible returned callable provenance to +// the existing contextual call proof. It does not attest closure ownership or +// purity. All executable returns must resolve to the same callable instance. +func ps6136Call(context *ps6125SSAContext, call *ssa.Call) *ps6125SSAContext { + if child := context.call(call); child != nil { + return child + } + if context == nil || call == nil || call.Parent() != context.flow.function || !context.flow.blocks[call.Block()] || call.Call.IsInvoke() || context.budget.remaining == 0 { + return nil + } + origin := ps6136ReturnedCallable(context.reference(call.Call.Value), make(map[ps6125SSAReference]bool)) + return ps6136CallOrigin(context, call, origin) +} + +// ps6136CallOrigin requires an already source-proved callable identity. It is +// not an opaque-function configuration escape hatch: immutable owner-field +// binding callers must close the source field's complete observation inventory. +func ps6136CallOrigin(context *ps6125SSAContext, call *ssa.Call, origin ps6125SSAReference) *ps6125SSAContext { + if context == nil || call == nil || origin.context == nil || call.Parent() != context.flow.function || !context.flow.blocks[call.Block()] || call.Call.IsInvoke() || context.budget.remaining == 0 { + return nil + } + function, captures := ps6125SSACallee(origin.value) + if function == nil || len(function.Params) != len(call.Call.Args) || len(function.FreeVars) != len(captures) || !types.Identical(function.Signature, call.Call.Value.Type()) { + return nil + } + for ancestor := context; ancestor != nil; ancestor = ancestor.parent { + if ancestor.flow.function == function { + return nil + } + } + inputs := make(map[*ssa.Parameter]ps6125Scalar, len(function.Params)) + lengths := make(map[*ssa.Parameter]ps6125Extent, len(function.Params)) + bindings := make(map[ssa.Value]ps6125SSAReference, len(function.Params)+len(function.FreeVars)) + for index, parameter := range function.Params { + argument := call.Call.Args[index] + bindings[parameter] = context.reference(argument) + inputs[parameter] = context.scalar(argument) + lengths[parameter] = context.length(argument) + } + for index, capture := range function.FreeVars { + bindings[capture] = origin.context.reference(captures[index]) + } + context.budget.remaining-- + child := &ps6125SSAContext{flow: ps6125AnalyzeSSAExtents(function, inputs, lengths), parent: context, site: call, budget: context.budget, bindings: bindings} + if context.calls == nil { + context.calls = make(map[*ssa.Call]*ps6125SSAContext) + } + context.calls[call] = child + return child +} + +func ps6136ReturnedCallable(reference ps6125SSAReference, active map[ps6125SSAReference]bool) ps6125SSAReference { + if reference.value == nil || reference.context == nil || active[reference] { + return ps6125SSAReference{} + } + if known := reference.context.callable(reference.value); known.value != nil { + return known + } + active[reference] = true + defer delete(active, reference) + if load, ok := reference.value.(*ssa.UnOp); ok && load.Op == token.MUL { + cell := reference.context.reference(load.X) + allocation, ok := cell.value.(*ssa.Alloc) + if !ok || cell.context == nil { + return ps6125SSAReference{} + } + store := cell.context.functionCell(allocation) + if store == nil { + return ps6125SSAReference{} + } + var point ssa.Instruction = load + current := reference.context + for current != cell.context && current.parent != nil { + point = current.site + current = current.parent + } + if current != cell.context || !cell.context.flow.instructionDominates(store, point) { + return ps6125SSAReference{} + } + return ps6136ReturnedCallable(cell.context.reference(store.Val), active) + } + call, ok := reference.value.(*ssa.Call) + if !ok { + return ps6125SSAReference{} + } + child := ps6136Call(reference.context, call) + if child == nil { + return ps6125SSAReference{} + } + var result ps6125SSAReference + for _, block := range child.flow.function.Blocks { + if !child.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + returned, ok := instruction.(*ssa.Return) + if !ok { + continue + } + if len(returned.Results) != 1 { + return ps6125SSAReference{} + } + value := ps6136ReturnedCallable(child.reference(returned.Results[0]), active) + if value.value == nil || result.value != nil && result != value { + return ps6125SSAReference{} + } + result = value + } + } + return result +} diff --git a/checks/ps6136_returned_callable_test.go b/checks/ps6136_returned_callable_test.go new file mode 100644 index 00000000..c235545f --- /dev/null +++ b/checks/ps6136_returned_callable_test.go @@ -0,0 +1,45 @@ +package checks + +import "testing" + +func TestPS6136ReturnedCallableBoundaries(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name, maker, body string + known bool + }{ + {"closed_return", `return func(){leaf()}`, `mk:=makeFactory();mk()`, true}, + {"returned_named", `return leaf`, `mk:=makeFactory();mk()`, true}, + {"unknown_dispatch", `if flag{return leaf};return other`, `mk:=makeFactory();mk()`, false}, + {"nil_dispatch", `if flag{return leaf};return nil`, `mk:=makeFactory();mk()`, false}, + {"rebound_cell", `return leaf`, `mk:=makeFactory();mk=other;mk()`, true}, + {"escaped_cell", `return leaf`, `mk:=makeFactory();expose(&mk);mk()`, false}, + {"captured_write", `return leaf`, `mk:=makeFactory();change:=func(){mk=other};change();mk()`, false}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + pkg := ps6125TestSSA(t, `package extent;var flag bool;func leaf(){};func other(){};func expose(*func()){};func makeFactory()func(){`+test.maker+`};func root(){`+test.body+`}`) + root := ps6125NewSSAContext(pkg.Func("root"), nil, nil, 24) + found, known := false, false + for _, call := range ps6125ContextCalls(root) { + if call.Call.StaticCallee() != nil { + continue + } + found = true + known = ps6136Call(root, call) != nil + } + // A compiler-proven direct replacement dispatch is accepted only + // as that actual replacement, never as the original allocator. + if !found && test.name == "rebound_cell" { + for _, call := range ps6125ContextCalls(root) { + if call.Call.StaticCallee() == pkg.Func("other") { + found, known = true, true + } + } + } + if !found || known != test.known { + t.Fatalf("dispatch found=%v closed=%v, want %v", found, known, test.known) + } + }) + } +} diff --git a/checks/ps6136_selection.go b/checks/ps6136_selection.go new file mode 100644 index 00000000..87295375 --- /dev/null +++ b/checks/ps6136_selection.go @@ -0,0 +1,112 @@ +package checks + +import ( + "go/types" + + "github.com/jxsl13/perfscan/config" + "golang.org/x/tools/go/ssa" +) + +func (context *ps6136ContractContext) selection(pkg *ssa.Package, entry *types.Func, contract *config.OutputWorkspaceContract) *ps6136Selection { + owner, model := context.named(contract.OwnerType), context.named(contract.ModelType) + factory := context.callable(contract.FactoryFunction) + if owner == nil || model == nil || entry == nil || factory == nil || owner.Obj().Pkg() != context.pass.Pkg || context.declarations[entry] == nil || context.declarations[factory] == nil { + return nil + } + function := pkg.Prog.FuncValue(entry) + constructor := pkg.Prog.FuncValue(factory) + if function == nil || constructor == nil || function.Signature.Params().Len() == 0 || constructor.Signature.Params().Len() == 0 || !types.Identical(function.Signature.Params().At(0).Type(), types.NewPointer(model)) || !types.Identical(constructor.Signature.Params().At(0).Type(), types.NewPointer(model)) || constructor.Signature.Results().Len() != 2 || !types.Identical(constructor.Signature.Results().At(0).Type(), types.NewPointer(owner)) || !types.Identical(constructor.Signature.Results().At(1).Type(), types.Universe.Lookup("error").Type()) { + return nil + } + root := ps6125NewSSAContext(function, nil, nil, 16384) + var selected *ps6125SSAContext + if function == constructor { + selected = root + } else { + remaining := 16384 + valid := ps6136WalkConsumerCalls(root, func(current *ps6125SSAContext, call *ssa.Call) bool { + if call.Call.StaticCallee() == constructor { + if selected != nil { + selected = nil + remaining = 0 + return true + } + selected = ps6136Call(current, call) + return true + } + return false + }, &remaining) + if !valid { + return nil + } + } + if selected == nil { + return nil + } + selection := &ps6136Selection{context: selected, ownerType: owner, modelType: model, model: selected.reference(constructor.Params[0]), allocationFunction: contract.AllocationFunction, releaseMethod: contract.ReleaseMethod, shapeMethod: contract.HeadShapeMethod, allocatorIDs: make(map[string]bool)} + selection.maximumRows, selection.width, selection.workspace = ps6136FieldVar(owner, contract.MaximumRowsField), ps6136FieldVar(owner, contract.WidthField), ps6136FieldVar(owner, contract.WorkspaceField) + selection.ops, selection.retained, selection.projector = ps6136FieldVar(owner, contract.BackendOpsField), ps6136FieldVar(owner, contract.RetainedListField), ps6136FieldVar(owner, contract.ProjectorField) + selection.modelConfig, selection.modelHead = ps6136FieldVar(model, contract.ModelConfigField), ps6136FieldVar(model, contract.ModelHeadField) + if selection.maximumRows == nil || selection.width == nil || selection.workspace == nil || selection.ops == nil || selection.retained == nil || selection.projector == nil || selection.modelConfig == nil || selection.modelHead == nil || !types.Identical(selection.maximumRows.Type(), types.Typ[types.Int]) || !types.Identical(selection.width.Type(), types.Typ[types.Int]) { + return nil + } + selection.configRows, selection.configWidth = ps6136FieldVar(selection.modelConfig.Type(), contract.ConfigMaximumRowsField), ps6136FieldVar(selection.modelConfig.Type(), contract.ConfigWidthField) + selection.retainedReleaseMethod = contract.RetainedBufferReleaseMethod + selection.slot = ps6136FieldVar(selection.workspace.Type(), contract.SlotBufferField) + selection.allocator = ps6136FieldVar(selection.ops.Type(), contract.BackendAllocatorField) + selection.primaryRecorder = ps6136FieldVar(selection.ops.Type(), contract.BackendRecorderField) + if contract.SecondaryRecorderField != "" { + selection.secondaryRecorder = ps6136FieldVar(selection.ops.Type(), contract.SecondaryRecorderField) + } + projector := context.named(contract.ProjectorType) + if selection.configRows == nil || selection.configWidth == nil || selection.slot == nil || selection.allocator == nil || projector == nil || !types.Identical(selection.configRows.Type(), types.Typ[types.Int]) || !types.Identical(selection.configWidth.Type(), types.Typ[types.Int]) { + return nil + } + projectorWidth := ps6136FieldVar(projector, contract.ProjectorWidthField) + if projectorWidth == nil { + return nil + } + selection.projectorWidths = map[*types.Named]*types.Var{projector: projectorWidth} + structure, ok := selection.ops.Type().Underlying().(*types.Struct) + if !ok { + return nil + } + selection.allocatorIndex = -1 + for index := 0; index < structure.NumFields(); index++ { + if structure.Field(index) == selection.allocator { + selection.allocatorIndex = index + } + } + if selection.allocatorIndex < 0 { + return nil + } + for _, identity := range contract.BackendAllocators { + if context.callable(identity) == nil { + return nil + } + selection.allocatorIDs[identity] = true + } + for _, block := range constructor.Blocks { + if !selected.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + returned, ok := instruction.(*ssa.Return) + if !ok || len(returned.Results) == 0 { + continue + } + if value, ok := returned.Results[0].(*ssa.Const); ok && value.IsNil() { + continue + } + identity := ps6136OwnerRoot(selected, returned.Results[0], owner) + if _, fresh := identity.value.(*ssa.Alloc); !fresh || selection.owner.value != nil && selection.owner != identity { + return nil + } + selection.owner = identity + } + } + if selection.owner.value == nil { + return nil + } + return selection +} diff --git a/checks/ps6136_snapshot.go b/checks/ps6136_snapshot.go new file mode 100644 index 00000000..17fcb6fd --- /dev/null +++ b/checks/ps6136_snapshot.go @@ -0,0 +1,66 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136InitialOwnerField returns the one source initializer of a fresh owner +// field. This alone is NOT immutable memory proof: the selected constructor's +// full field-use/owner-alias inventory must independently establish that no +// escaped/captured/other-method write can change this field before its consumer. +func ps6136InitialOwnerField(owner ps6125SSAReference, field *types.Var) ps6125SSAReference { + allocation, ok := owner.value.(*ssa.Alloc) + if !ok || owner.context == nil || field == nil { + return ps6125SSAReference{} + } + var source *ssa.Store + for _, block := range owner.context.flow.function.Blocks { + if !owner.context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + address, ok := instruction.(*ssa.FieldAddr) + if !ok || address.X != allocation { + continue + } + pointer, ok := allocation.Type().Underlying().(*types.Pointer) + if !ok { + return ps6125SSAReference{} + } + structure, ok := pointer.Elem().Underlying().(*types.Struct) + if !ok || structure.Field(address.Field) != field { + continue + } + users := address.Referrers() + if users == nil { + return ps6125SSAReference{} + } + for _, user := range *users { + if user.Parent() == owner.context.flow.function && !owner.context.flow.blocks[user.Block()] { + continue + } + switch use := user.(type) { + case *ssa.DebugRef: + case *ssa.UnOp: + if use.Op != token.MUL || use.X != address { + return ps6125SSAReference{} + } + case *ssa.Store: + if use.Addr != address || source != nil { + return ps6125SSAReference{} + } + source = use + default: + return ps6125SSAReference{} + } + } + } + } + if source == nil { + return ps6125SSAReference{} + } + return owner.context.reference(source.Val) +} diff --git a/checks/ps6136_snapshot_test.go b/checks/ps6136_snapshot_test.go new file mode 100644 index 00000000..38276a74 --- /dev/null +++ b/checks/ps6136_snapshot_test.go @@ -0,0 +1,79 @@ +package checks + +import ( + "go/types" + "testing" + + "golang.org/x/tools/go/ssa" +) + +func TestPS6136AuthenticSharedPublicAllocatorSnapshot(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + t.Run(revision, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, revision, true, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + entry := ps6125NewSSAContext(pkg.Func("New"), nil, nil, 128) + var root *ps6125SSAContext + for _, call := range ps6125ContextCalls(entry) { + if call.Call.StaticCallee() == pkg.Func("newDecoder") { + root = entry.call(call) + } + } + ownerType := pkg.Pkg.Scope().Lookup("Decoder").Type().(*types.Named) + workspace, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, "logits") + ops, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, "ops") + backendType := pkg.Pkg.Scope().Lookup("backendOps").Type().Underlying().(*types.Struct) + allocator, indices, _ := types.LookupFieldOrMethod(pkg.Pkg.Scope().Lookup("backendOps").Type(), true, pkg.Pkg, "newBuffer") + slot := pkg.Pkg.Scope().Lookup("bufSlot").Type().Underlying().(*types.Struct).Field(0) + found := 0 + for _, call := range ps6125ContextCalls(root) { + if call.Call.StaticCallee() == nil || call.Call.StaticCallee().Name() != "allocScratch" { + continue + } + allocation := root.call(call) + paths := ps6125AccessPaths{flow: allocation.flow} + for _, block := range allocation.flow.function.Blocks { + for _, instruction := range block.Instrs { + write, ok := instruction.(*ssa.Store) + if !ok { + continue + } + path := paths.resolve(write.Addr) + if !path.known || len(path.access.fields) != 1 || path.access.fields[0] != workspace { + continue + } + factoryCall := write.Val.(*ssa.Call) + factory := ps6136Call(allocation, factoryCall) + input := allocation.reference(factoryCall.Call.Args[0]) + backend := ps6136FactoryResult(factory, input, allocator.(*types.Var), slot) + if cell := ps6136FactoryErrorCell(factory, backend); cell.value == nil || cell.context != root { + t.Fatal("returned factory allocator failure does not feed the same selected constructor error cell") + } + identity := ps6136AccessOwnerRoot(allocation, path, ownerType) + method, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, "Release") + if !ps6136ConstructorErrorBarrier(root, identity, ownerType, ps6136FactoryErrorCell(factory, backend), ps6090FunctionID(method.(*types.Func))) { + t.Fatal("actual shared same-owner failure release/publication barrier unproved") + } + snapshot := ps6136InitialOwnerField(identity, ops.(*types.Var)) + if backend == nil || snapshot.value == nil || backendType.Field(indices[0]) != allocator { + t.Fatal("actual public constructor ops initializer unavailable") + } + origin := snapshot.context.structField(snapshot.value, indices[0]) + if ps6136CallOrigin(factory, backend, origin) == nil || !ps6136BoundAllocator(factory, backend, input, map[string]bool{"github.com/jxsl13/goai/backend/metal.NewDeviceBufferF32": true}) { + t.Fatal("actual shared public constructor -> initial ops -> captured factory -> exact native allocator binding unproved") + } + found++ + } + } + } + if found != 1 { + t.Fatalf("public shared output allocations %d", found) + } + }) + } +} diff --git a/checks/ps6136_source_functions.go b/checks/ps6136_source_functions.go new file mode 100644 index 00000000..50f43153 --- /dev/null +++ b/checks/ps6136_source_functions.go @@ -0,0 +1,43 @@ +package checks + +import ( + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// Include methods on every local named type, not only the selected owner: +// wrapper/helper methods may introduce observations of that owner's fields. +func ps6136SourceFunctions(pkg *ssa.Package) []*ssa.Function { + if pkg == nil { + return nil + } + seen := make(map[*ssa.Function]bool) + var result []*ssa.Function + var add func(*ssa.Function) + add = func(function *ssa.Function) { + if function == nil || seen[function] { + return + } + seen[function] = true + result = append(result, function) + for _, child := range function.AnonFuncs { + add(child) + } + } + for _, member := range pkg.Members { + switch member := member.(type) { + case *ssa.Function: + add(member) + case *ssa.Type: + named, ok := types.Unalias(member.Type()).(*types.Named) + if !ok || named.Obj().Pkg() != pkg.Pkg { + continue + } + for index := 0; index < named.NumMethods(); index++ { + add(pkg.Prog.FuncValue(named.Method(index))) + } + } + } + return result +} diff --git a/checks/ps6136_struct_zero.go b/checks/ps6136_struct_zero.go new file mode 100644 index 00000000..ea4424b7 --- /dev/null +++ b/checks/ps6136_struct_zero.go @@ -0,0 +1,108 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136StructFieldZero proves a nil function or false bool in a closed source +// value struct, including an omitted literal field. It never treats an unknown +// heap field, unavailable struct projection or opaque helper as zero. Subsequent +// owner-memory invariance remains an independent full-observation obligation. +func ps6136StructFieldZero(context *ps6125SSAContext, value ssa.Value, field int, budget int) bool { + if context == nil || value == nil || budget <= 0 { + return false + } + reference := context.reference(value) + if reference.value == nil { + return false + } + if reference.context != context || reference.value != value { + return ps6136StructFieldZero(reference.context, reference.value, field, budget-1) + } + structure, ok := value.Type().Underlying().(*types.Struct) + if !ok || field < 0 || field >= structure.NumFields() { + return false + } + kind := structure.Field(field).Type().Underlying() + _, function := kind.(*types.Signature) + boolean := types.Identical(structure.Field(field).Type(), types.Typ[types.Bool]) + if !function && !boolean { + return false + } + // SSA Const's nil Value represents the build-time zero of any type, + // including an entire source value struct (not an unknown runtime load). + if literal, ok := value.(*ssa.Const); ok && literal.Value == nil { + return true + } + if projected := context.structField(value, field); projected.value != nil { + literal, ok := projected.value.(*ssa.Const) + if !ok { + return false + } + if function { + return literal.IsNil() + } + fact := projected.context.scalar(literal) + return fact.state == ps6125Boolean && !fact.truth + } + load, ok := value.(*ssa.UnOp) + if !ok || load.Op != token.MUL { + return false + } + return ps6136StructFieldZeroAt(context, load.X, field, load, budget-1) +} + +func ps6136StructFieldZeroAt(context *ps6125SSAContext, address ssa.Value, field int, point ssa.Instruction, budget int) bool { + if context == nil || address == nil || point == nil || budget <= 0 { + return false + } + cell := context.reference(address) + allocation, ok := cell.value.(*ssa.Alloc) + if !ok { + return false + } + pointer, ok := allocation.Type().Underlying().(*types.Pointer) + if !ok { + return false + } + structure, ok := pointer.Elem().Underlying().(*types.Struct) + if !ok || field < 0 || field >= structure.NumFields() { + return false + } + _, function := structure.Field(field).Type().Underlying().(*types.Signature) + if !function && !types.Identical(structure.Field(field).Type(), types.Typ[types.Bool]) { + return false + } + stores := cell.context.structCell(allocation) + if stores == nil || field < 0 || field >= len(stores.fields) { + return false + } + // Use the actual read/call edge, not lexical closure creation. A load of + // another invocation's cell cannot inherit zero facts from its same type. + current := context + for current != cell.context && current.parent != nil { + point = current.site + current = current.parent + } + if current != cell.context || !cell.context.flow.instructionDominates(allocation, point) { + return false + } + if stores.whole != nil { + return cell.context.flow.instructionDominates(stores.whole, point) && ps6136StructFieldZero(cell.context, stores.whole.Val, field, budget-1) + } + if store := stores.fields[field]; store != nil { + literal, ok := store.Val.(*ssa.Const) + if !ok || !cell.context.flow.instructionDominates(store, point) { + return false + } + if _, function := literal.Type().Underlying().(*types.Signature); function { + return literal.IsNil() + } + fact := cell.context.scalar(literal) + return types.Identical(literal.Type(), types.Typ[types.Bool]) && fact.state == ps6125Boolean && !fact.truth + } + return true +} diff --git a/checks/ps6136_struct_zero_test.go b/checks/ps6136_struct_zero_test.go new file mode 100644 index 00000000..9a43075f --- /dev/null +++ b/checks/ps6136_struct_zero_test.go @@ -0,0 +1,41 @@ +package checks + +import "testing" + +func TestPS6136ClosedStructZero(t *testing.T) { + t.Parallel() + pkg := ps6125TestSSA(t, `package extent +type ops struct{fn func(); flag bool; other int} +func leaf(ops){} +func escape(*ops){} +func zero(){o:=ops{other:1};leaf(o)} +func explicit(){o:=ops{fn:nil,flag:false};leaf(o)} +func nonzero(){o:=ops{fn:func(){},flag:true};leaf(o)} +func borrowed(o ops){leaf(o)} +func mixed(){o:=ops{};o.other=1;leaf(o)} +func escaped(){o:=ops{};escape(&o);leaf(o)} +func late(){o:=ops{};leaf(o);o.fn=func(){}} +func captured(){o:=ops{};fn:=func(){o.flag=true};fn();leaf(o)} +`) + for _, name := range []string{"zero", "explicit", "nonzero", "borrowed", "mixed", "escaped", "late", "captured"} { + t.Run(name, func(t *testing.T) { + t.Parallel() + context := ps6125NewSSAContext(pkg.Func(name), nil, nil, 64) + count := 0 + for _, call := range ps6125ContextCalls(context) { + if call.Call.StaticCallee() != pkg.Func("leaf") { + continue + } + for _, field := range []int{0, 1} { + if got := ps6136StructFieldZero(context, call.Call.Args[0], field, 32); got != (name == "zero" || name == "explicit") { + t.Fatalf("%s field %d zero = %v", name, field, got) + } + } + count++ + } + if count != 1 { + t.Fatalf("leaf count %d", count) + } + }) + } +} diff --git a/checks/ps6136_test.go b/checks/ps6136_test.go new file mode 100644 index 00000000..e152f67b --- /dev/null +++ b/checks/ps6136_test.go @@ -0,0 +1,229 @@ +package checks + +import ( + "go/ast" + "go/types" + "testing" + + "github.com/jxsl13/perfscan/config" + "golang.org/x/tools/go/analysis" + "golang.org/x/tools/go/ssa" +) + +func TestPS6136SelectedOwnerGeometryInitialization(t *testing.T) { + t.Parallel() + for _, revision := range []string{"before", "after"} { + for _, entryName := range []string{"NewGPT", "New"} { + t.Run(revision+"/"+entryName, func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwners(t, revision, true, true) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + entry := ps6125NewSSAContext(pkg.Func(entryName), nil, nil, 16384) + ownerName, constructorName := "GPTDecoder", "newGPTDecoder" + if entryName == "New" { + ownerName, constructorName = "Decoder", "newDecoder" + } + function := pkg.Func(constructorName) + var context *ps6125SSAContext + for _, call := range ps6125ContextCalls(entry) { + if call.Call.StaticCallee() == function { + context = entry.call(call) + } + } + if context == nil { + t.Fatal("selected constructor missing") + } + ownerType := pkg.Pkg.Scope().Lookup(ownerName).Type().(*types.Named) + modelType := function.Params[0].Type().Underlying().(*types.Pointer).Elem().(*types.Named) + field := func(typ types.Type, name string) *types.Var { + object, _, _ := types.LookupFieldOrMethod(typ, true, pkg.Pkg, name) + return object.(*types.Var) + } + selection := &ps6136Selection{context: context, ownerType: ownerType, modelType: modelType, model: context.reference(function.Params[0]), maximumRows: field(ownerType, "maxLen"), width: field(ownerType, "v"), modelConfig: field(modelType, "Config")} + selection.configRows, selection.configWidth = field(selection.modelConfig.Type(), "Ctx"), field(selection.modelConfig.Type(), "Vocab") + for _, block := range function.Blocks { + for _, instruction := range block.Instrs { + returned, ok := instruction.(*ssa.Return) + if !ok || len(returned.Results) == 0 { + continue + } + if constant, ok := returned.Results[0].(*ssa.Const); ok && constant.IsNil() { + continue + } + selection.owner = ps6136OwnerRoot(context, returned.Results[0], ownerType) + } + } + if !selection.initializationValues(16384) { + t.Fatal("actual selected model config -> same fresh owner's maximum/width source flow missing") + } + selection.workspace, selection.ops, selection.retained = field(ownerType, "logits"), field(ownerType, "ops"), field(ownerType, "all") + selection.primaryRecorder = field(selection.ops.Type(), "newRecorder") + selection.secondaryRecorder = field(selection.ops.Type(), "newDecodeRecorder") + projectorName, headName := "head", "Head" + allocationFunction := function + if ownerName == "Decoder" { + projectorName, headName = "out", "Out" + object, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, "allocScratch") + allocationFunction = pkg.Prog.FuncValue(object.(*types.Func)) + } + var consumers []*ps6136ConsumerProof + selection.projector, selection.modelHead = field(ownerType, projectorName), field(modelType, headName) + selection.allocationFunction = ps6090FunctionID(allocationFunction.Object().(*types.Func)) + selection.slot = pkg.Pkg.Scope().Lookup("bufSlot").Type().Underlying().(*types.Struct).Field(0) + allocator, index, _ := types.LookupFieldOrMethod(selection.ops.Type(), true, pkg.Pkg, "newBuffer") + selection.allocator, selection.allocatorIndex = allocator.(*types.Var), index[0] + selection.allocatorIDs = map[string]bool{"github.com/jxsl13/goai/backend/metal.NewDeviceBufferF32": true} + release, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, "Release") + selection.releaseMethod = ps6090FunctionID(release.(*types.Func)) + selection.retainedReleaseMethod = "github.com/jxsl13/goai/llamagpu.buffer.Release" + shape, _, _ := types.LookupFieldOrMethod(selection.modelHead.Type(), true, pkg.Pkg, "Shape") + selection.shapeMethod = ps6090FunctionID(shape.(*types.Func)) + projectorType := pkg.Pkg.Scope().Lookup("f32Linear").Type().(*types.Named) + selection.projectorWidths = map[*types.Named]*types.Var{projectorType: field(projectorType, "n")} + proof := selection.constructorAllocation(16384) + if (proof != nil) != (revision == "before") { + t.Fatalf("selected original constructor allocation proof=%v", proof != nil) + } + if proof != nil && !selection.constructorLifetime(proof, 16384) { + t.Fatal("actual selected constructor allocator/error-cell/earlier-return lifetime closure missing") + } + if proof != nil { + backend := ps6136FactoryResult(proof.factory, proof.factory.reference(proof.factory.flow.function.Params[0]), selection.allocator, selection.slot) + wrapper := pkg.Pkg.Scope().Lookup("mBuf").Type().(*types.Named) + if backend == nil || !ps6136BoundBufferWrapper(proof.factory, backend, wrapper, wrapper.Underlying().(*types.Struct).Field(0)) { + t.Fatal("exact allocator wrapper/native field is not joined to buffer bridge") + } + other := types.NewNamed(types.NewTypeName(0, pkg.Pkg, "otherBufferWrapper", nil), wrapper.Underlying(), nil) + if ps6136BoundBufferWrapper(proof.factory, backend, other, wrapper.Underlying().(*types.Struct).Field(0)) { + t.Fatal("different same-native-field wrapper accepted") + } + } + if revision == "after" { + selection.configRows, selection.configWidth = selection.configWidth, selection.configRows + if selection.initializationValues(16384) { + t.Fatal("swapped source geometry identity accepted") + } + return + } + leaves := ps6136OwnerLeaves(pkg) + methods := []string{"Step", "StepNLast", "StepN", "Generate"} + if ownerName == "Decoder" { + methods = append(methods, "ProfileMetalStep") + } + for _, name := range methods { + object, _, _ := types.LookupFieldOrMethod(ownerType, true, pkg.Pkg, name) + consumer := pkg.Prog.FuncValue(object.(*types.Func)) + consumerContext := ps6125NewSSAContext(consumer, nil, nil, 16384) + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg} + got := selection.consumerCalls(consumerContext, leaves, name == "StepN", func(current *ps6125SSAContext, call *ssa.Call, leaf *config.OutputWorkspaceLeaf) bool { + return ps6136FixtureCapacityPrefix(pass, pkg, current, call, selection.width) + }) + if got == nil { + t.Fatalf("%s complete selected consumer call/read proof missing", name) + } + consumers = append(consumers, got) + if ownerName == "Decoder" && name == "Generate" && got.capacityTransfers != 1 { + t.Fatalf("physical transfer census %d", got.capacityTransfers) + } + } + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg} + recorderWrapper := pkg.Pkg.Scope().Lookup("mRec").Type().(*types.Named) + if !selection.recorderBinding(pass, field(selection.ops.Type(), "newRecorder"), "github.com/jxsl13/goai/backend/metal.NewRecorder", recorderWrapper, field(recorderWrapper, "r")) { + t.Fatal("selected constructor's exact typed native recorder factory binding missing") + } + contract := ps6136NativeOwnerContract(pkg, ownerName) + if !ps6136ContractsContext(pass).nativeBindings(selection, proof, &contract) { + t.Fatal("complete selected native factory/adapter/capability/state closure missing") + } + if !selection.closedSource(pass, pkg, consumers) { + t.Fatal("complete selected-instance constructor/field-use/noninterference closure missing") + } + selection.configRows, selection.configWidth = selection.configWidth, selection.configRows + if selection.initializationValues(16384) { + t.Fatal("swapped source geometry identity accepted") + } + }) + } + } +} + +func ps6136OwnerLeaves(pkg *ssa.Package) []config.OutputWorkspaceLeaf { + leaf := func(typ, name, kind string, roles []int, rows, width, host int) config.OutputWorkspaceLeaf { + object, _, _ := types.LookupFieldOrMethod(pkg.Pkg.Scope().Lookup(typ).Type(), true, pkg.Pkg, name) + return config.OutputWorkspaceLeaf{Method: ps6090FunctionID(object.(*types.Func)), Kind: kind, BufferArguments: roles, RowsArgument: rows, WidthArgument: width, HostArgument: host, SemanticsReviewed: true} + } + return []config.OutputWorkspaceLeaf{ + leaf("linear", "record", "projection", []int{2}, 3, -1, -1), + leaf("recorder", "AddBias", "row-width", []int{0, 2}, 3, 4, -1), + leaf("buffer", "DownloadF32", "download", []int{-1}, -1, -1, 0), + leaf("deviceTopKer", "TopKN", "prefix", []int{-1}, -1, 0, -1), + leaf("deviceTopPer", "TopKN", "prefix", []int{-1}, -1, 0, -1), + leaf("deviceTopPer", "SoftmaxStatsN", "prefix", []int{-1}, -1, 0, -1), + leaf("deviceTopPer", "ToHost", "capacity-transfer", []int{-1}, -1, -1, -1), + } +} + +func ps6136NativeOwnerContract(pkg *ssa.Package, owner string) config.OutputWorkspaceContract { + local, native := "github.com/jxsl13/goai/llamagpu.", "github.com/jxsl13/goai/backend/metal." + contract := config.OutputWorkspaceContract{OwnerType: local + owner, BackendRecorderField: "newRecorder", SecondaryRecorderField: "newDecodeRecorder", NativeRecorderType: native + "Recorder", RecorderWrapperType: local + "mRec", RecorderNativeField: "r", NativeRecorderFactory: native + "NewRecorder", BufferWrapperType: local + "mBuf", BufferNativeField: "DeviceBuffer", BufferBridge: local + "mb", ProjectorType: local + "f32Linear", ProjectorRecordMethod: local + "f32Linear.record", ProjectorWeightField: "w", ProjectorInnerField: "k", ProjectorWidthField: "n", RecorderProjectionMethod: local + "recorder.MatMul", NativeProjectionMethod: native + "Recorder.MatMul", Leaves: ps6136OwnerLeaves(pkg)} + for _, name := range []string{"Step", "StepNLast", "Generate"} { + contract.CommonMethods = append(contract.CommonMethods, local+owner+"."+name) + } + contract.BulkMethod = local + owner + ".StepN" + if owner == "Decoder" { + contract.RecorderOverrideMethod = local + "Decoder.ProfileMetalStep" + contract.DropPendingMethod = local + "Decoder.dropPending" + contract.NativeOverrideFactory = native + "NewProfilingRecorder" + contract.RecorderOverrideWrapperType = local + "mProfileRec" + contract.RecorderOverrideEmbedField = "mRec" + contract.RecorderOverrideMetadataField = "profile" + contract.AsyncRecorderField = "asyncEncode" + contract.CommonMethods = append(contract.CommonMethods, contract.RecorderOverrideMethod) + } + for index := range contract.Leaves { + leaf := &contract.Leaves[index] + switch index { + case 0: + leaf.Implementations = []string{local + "f32Linear.record"} + case 1: + leaf.Implementations = []string{native + "Recorder.AddBias"} + case 2: + leaf.Implementations = []string{native + "DeviceBuffer.DownloadF32"} + case 3: + leaf.Implementations = []string{native + "DeviceBuffer.TopKN"} + default: + leaf.SelectedCapabilityAbsenceAllowed = true + } + } + return contract +} + +func ps6136FixtureCapacityPrefix(pass *analysis.Pass, pkg *ssa.Package, current *ps6125SSAContext, call *ssa.Call, width *types.Var) bool { + var tensorPackage *types.Package + for _, imported := range pkg.Pkg.Imports() { + if imported.Path() == "github.com/jxsl13/goai/tensor" { + tensorPackage = imported + } + } + if tensorPackage == nil || len(current.flow.function.Params) == 0 { + return false + } + storage, _, _ := types.LookupFieldOrMethod(types.NewPointer(tensorPackage.Scope().Lookup("Tensor").Type()), true, pkg.Pkg, "Storage") + f32, _, _ := types.LookupFieldOrMethod(types.NewPointer(tensorPackage.Scope().Lookup("Storage").Type()), true, pkg.Pkg, "F32") + var transfer *ast.CallExpr + ast.Inspect(current.flow.function.Syntax(), func(node ast.Node) bool { + candidate, ok := node.(*ast.CallExpr) + if !ok || call.Pos() < candidate.Pos() || candidate.End() <= call.Pos() { + return true + } + selector, ok := candidate.Fun.(*ast.SelectorExpr) + if ok && pass.TypesInfo.Selections[selector] != nil && pass.TypesInfo.Selections[selector].Obj() == call.Call.Method { + transfer = candidate + } + return true + }) + return transfer != nil && ps6136HostPrefix(pass, transfer, current.flow.function.Params[0].Object(), width, storage.(*types.Func), f32.(*types.Func)) +} diff --git a/checks/ps6136_write_nodes.go b/checks/ps6136_write_nodes.go new file mode 100644 index 00000000..34f22cff --- /dev/null +++ b/checks/ps6136_write_nodes.go @@ -0,0 +1,89 @@ +package checks + +import ( + "go/ast" + "go/types" + + "golang.org/x/tools/go/analysis" + "golang.org/x/tools/go/ssa" +) + +// Map actual source store positions to exact initialization targets, never to +// a function/block-wide allowance that could conceal contradictory extra uses. +func ps6136ConstructorWriteNodes(pass *analysis.Pass, pkg *ssa.Package, owner *types.Named, fields map[*types.Var]bool, writes *ps6136ConstructorWrites) map[ast.Node]bool { + if pass == nil || writes == nil { + return nil + } + functions := ps6136SourceFunctions(pkg) + constructors := ps6136ConstructorInventory(pkg, owner) + if constructors == nil { + return nil + } + for writer := range writes.writers { + object, ok := writer.Object().(*types.Func) + if !ok || !object.Exported() || constructors[writer].value != nil { + continue + } + // Internal construction calls cannot establish that an exported + // setter is constructor-only; clients may invoke it on a published + // instance. Actual output/geometry initializing helpers are private. + return nil + } + proved := make(map[ast.Node]bool) + valid := true + for _, file := range pass.Files { + ast.Inspect(file, func(node ast.Node) bool { + if !valid { + return false + } + var target ast.Node + switch occurrence := node.(type) { + case *ast.SelectorExpr: + if selection := pass.TypesInfo.Selections[occurrence]; selection != nil { + if field, ok := selection.Obj().(*types.Var); ok && fields[field] { + target = occurrence + } + } + case *ast.KeyValueExpr: + if key, ok := occurrence.Key.(*ast.Ident); ok { + if field, ok := pass.TypesInfo.Uses[key].(*types.Var); ok && fields[field] { + target = occurrence + } + } + case *ast.CallExpr: + function, _, ok := typedCallee(pass, occurrence.Fun) + if !ok { + break + } + callee := pkg.Prog.FuncValue(function) + if !writes.writers[callee] || constructors[callee].value != nil { + break + } + // A constructor-only mutating helper cannot later be invoked + // by an inference/wrapper method on a published owner. + var caller *ssa.Function + for _, candidate := range functions { + syntax := candidate.Syntax() + if syntax != nil && syntax.Pos() <= occurrence.Pos() && occurrence.End() <= syntax.End() && (caller == nil || caller.Syntax().Pos() < syntax.Pos()) { + caller = candidate + } + } + if caller == nil || !writes.contexts[caller] { + valid = false + } + } + if target != nil { + for position := range writes.positions { + if target.Pos() <= position && position < target.End() { + proved[target] = true + } + } + } + return valid + }) + } + if !valid { + return nil + } + return proved +} diff --git a/checks/ps6136_zero_flow.go b/checks/ps6136_zero_flow.go new file mode 100644 index 00000000..e80b4893 --- /dev/null +++ b/checks/ps6136_zero_flow.go @@ -0,0 +1,87 @@ +package checks + +import ( + "go/token" + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// ps6136ZeroSourceFlow refines only CFG edges whose condition is established +// by a closed source struct's nil function/false bool. Unknown heap loads, +// pointer escapes and unresolved callbacks retain both branches. It does not +// establish immutable owner fields or source-call effects. +func ps6136ZeroSourceFlow(context *ps6125SSAContext) *ps6125SSAContext { + if context == nil || context.flow == nil || len(context.flow.function.Blocks) == 0 { + return context + } + flow := *context.flow + flow.blocks = make(map[*ssa.BasicBlock]bool) + flow.edges = make(map[ps6125SSAEdge]bool) + flow.origins = nil + clone := &ps6125SSAContext{flow: &flow, parent: context.parent, site: context.site, budget: context.budget, bindings: context.bindings} + pending := []*ssa.BasicBlock{flow.function.Blocks[0]} + flow.edges[ps6125SSAEdge{to: flow.function.Blocks[0]}] = true + for len(pending) != 0 { + block := pending[len(pending)-1] + pending = pending[:len(pending)-1] + if flow.blocks[block] { + continue + } + flow.blocks[block] = true + successors := block.Succs + if len(block.Instrs) > 0 { + if condition, ok := block.Instrs[len(block.Instrs)-1].(*ssa.If); ok { + fact := context.scalar(condition.Cond) + truth, known := fact.truth, fact.state == ps6125Boolean + if !known { + truth, known = ps6136ZeroCondition(context, condition.Cond) + } + if known { + index := 1 + if truth { + index = 0 + } + successors = successors[index : index+1] + } + } + } + for _, next := range successors { + flow.edges[ps6125SSAEdge{from: block, to: next}] = true + pending = append(pending, next) + } + } + return clone +} + +func ps6136ZeroCondition(context *ps6125SSAContext, value ssa.Value) (bool, bool) { + if types.Identical(value.Type(), types.Typ[types.Bool]) && ps6136ZeroField(context, value) { + return false, true + } + comparison, ok := value.(*ssa.BinOp) + if !ok || comparison.Op != token.EQL && comparison.Op != token.NEQ { + return false, false + } + left, right := comparison.X, comparison.Y + if literal, ok := left.(*ssa.Const); ok && literal.IsNil() { + left, right = right, left + } + literal, ok := right.(*ssa.Const) + _, function := left.Type().Underlying().(*types.Signature) + if !ok || !literal.IsNil() || !function || !ps6136ZeroField(context, left) { + return false, false + } + return comparison.Op == token.EQL, true +} + +func ps6136ZeroField(context *ps6125SSAContext, value ssa.Value) bool { + if field, ok := value.(*ssa.Field); ok { + return ps6136StructFieldZero(context, field.X, field.Field, 64) + } + load, ok := value.(*ssa.UnOp) + if !ok || load.Op != token.MUL { + return false + } + field, ok := load.X.(*ssa.FieldAddr) + return ok && ps6136StructFieldZeroAt(context, field.X, field.Field, load, 64) +} diff --git a/checks/ps6136_zero_flow_test.go b/checks/ps6136_zero_flow_test.go new file mode 100644 index 00000000..fe0a909e --- /dev/null +++ b/checks/ps6136_zero_flow_test.go @@ -0,0 +1,60 @@ +package checks + +import "testing" + +func TestPS6136SourceZeroBranchSpecialization(t *testing.T) { + t.Parallel() + pkg := ps6125TestSSA(t, `package extent +type ops struct{fn func(); eager bool} +func small(){} +func large(){} +func mutate(*ops){} +func work(o ops){ + f:=func(){if o.fn==nil{small()}else{large()};if o.eager{large()}else{small()}} + f() +} +func zero(){work(ops{})} +func explicit(){work(ops{fn:nil,eager:false})} +func nonzero(){work(ops{fn:func(){},eager:true})} +func unknown(o ops){work(o)} +func changed(o ops){f:=func(){if o.fn==nil{small()}else{large()}};mutate(&o);f()} +func assignment(o ops){f:=func(){if o.fn==nil{small()}else{large()}};o.fn=func(){};f()} +`) + for _, name := range []string{"zero", "explicit", "nonzero", "unknown", "changed", "assignment"} { + t.Run(name, func(t *testing.T) { + t.Parallel() + context := ps6125NewSSAContext(pkg.Func(name), nil, nil, 64) + if name != "changed" && name != "assignment" { + context = context.call(ps6125ContextCalls(context)[0]) + } + var child *ps6125SSAContext + for _, call := range ps6125ContextCalls(context) { + if call.Call.StaticCallee() != pkg.Func("mutate") { + child = ps6136ZeroSourceFlow(ps6136Call(context, call)) + } + } + if child == nil { + t.Fatal("actual closed callback invocation unavailable") + } + small, large := 0, 0 + for _, call := range ps6125ContextCalls(child) { + if call.Call.StaticCallee() == pkg.Func("small") { + small++ + } + if call.Call.StaticCallee() == pkg.Func("large") { + large++ + } + } + wantSmall, wantLarge := 2, 2 + if name == "zero" || name == "explicit" { + wantLarge = 0 + } + if name == "changed" || name == "assignment" { + wantSmall, wantLarge = 1, 1 + } + if small != wantSmall || large != wantLarge { + t.Fatalf("source branches small=%d large=%d, want %d,%d", small, large, wantSmall, wantLarge) + } + }) + } +} diff --git a/checks/testdata/ps6136-owner/after/decoder.go.txt b/checks/testdata/ps6136-owner/after/decoder.go.txt new file mode 100644 index 00000000..ba8e5f36 --- /dev/null +++ b/checks/testdata/ps6136-owner/after/decoder.go.txt @@ -0,0 +1,3949 @@ +// Package llamagpu decodes nlp.Llama models on the GPU with batched command buffers (ADR-0019): +// each per-token step records the whole layer stack into ONE command buffer over device-resident +// weights and KV cache, instead of paying a dispatch round-trip per op. Measured on a real model +// (D=512, GQA 8:2, 6 layers): 24× faster than nlp.Llama.DecodeStep on Metal and 21× on Vulkan, +// with token-for-token identical greedy output (§T404/§T409). +// +// Architecture coverage (CUDA): the same batched Decoder core drives 20 model families beyond Llama, +// each a New*CUDA constructor + greedy parity test, built from one composable set of gated +// generalizations (norm type/placement — pre / one-norm-parallel / two-norm-parallel / post-norm / +// sandwich; RMSNorm / LayerNorm±bias / (1+w); MLP — SwiGLU / GELU / squared-ReLU / GeGLU; position — +// full RoPE / partial RoPE / ALiBi; attention width — MHA / GQA / MQA / decoupled-head_dim; per-head +// & full-width QK-norm; attention soft-cap; config scalars; logit scale; √dim embed; biased +// projections/lm_head). Dense: StableLM, StarCoder2, Phi, GPT-NeoX, Qwen2, Qwen3, Phi-3, Granite, +// Cohere, Nemotron, Gemma, OLMo2, Falcon, Gemma2, MPT. Sparse Mixture-of-Experts (route → experts → +// combine, all in the pre-recorded command buffer via on-device routing): Mixtral, Qwen3-MoE, OLMoE, +// GraniteMoE, Qwen2-MoE (shared expert). These New*CUDA constructors are cuda-only where partial +// rotary / soft-cap / ALiBi / MoE kernels are involved; metal/vulkan return the corresponding stubs. +// +// Build [New] (Metal, darwin+cgo) or [NewVulkan] (vulkan build tag) from any *nlp.Llama — including +// one loaded via nlp.LlamaFromGGUF — then call [Decoder.Generate] with any nlp.TokenSampler +// (nlp.Sampler or nlp.Mirostat), or drive +// [Decoder.Step] / [Decoder.StepN] yourself. Beyond plain generation the package provides: +// +// - [Decoder.StepN]: a whole multi-token window in one recorded step. Generate uses it to prefill +// the prompt in a single dispatch round-trip — measured 41× faster than token-at-a-time (§T418). +// - [NewQuant] / [NewQuantF16KV] / [NewQuantVulkan]: decode a quantized model +// (nlp.QuantizeLlama or a quantized GGUF) with every projection held device-resident in its +// 4-8× smaller ggml form (§T415). NewQuantF16KV additionally halves retained Metal K/V cache +// storage while preserving the existing f32-cache constructor as the default. +// - [SpeculativeGenerate]: lossless speculative decoding — a small draft Decoder proposes, the +// target verifies the window in one StepN, and the output stays exactly target-distributed +// (§T419; ~2-3× at typical acceptance rates with a trained draft). +// - [PromptLookupGenerate]: draft-model-free speculative decoding — candidate continuations are +// copied from the sequence's own history by n-gram matching (§T426); lossless, and effective +// when the output repeats the input (summarization, RAG, code editing). +// - [MedusaGenerate]: Medusa chain drafting (§T446/§T447) — nlp.MedusaHeads trained on the base's +// own rollouts draft host-side for free, one batched StepN verifies the window. Measured 1.81× +// at 97% acceptance where draft-model speculative managed 1.12× (drafting cost is the +// difference on dispatch-bound decoders); greedy-anchored, not distribution-exact. +// +// The Decoder core is backend-agnostic over small buffer/recorder interfaces; the two adapters plug +// in the concrete backends, whose exported Recorder/DeviceBuffer APIs are identical by construction +// (§T391/§T408). +package llamagpu + +import ( + "errors" + "fmt" + "math" + + "github.com/jxsl13/goai/backend" + "github.com/jxsl13/goai/nlp" + "github.com/jxsl13/goai/nn" + "github.com/jxsl13/goai/tensor" +) + +// unarySiLU/binaryAdd/binaryMul are the shared kernel selectors (identical on both backends — +// they must match shaders/unary.comp / metal_bridge.m's unary switch and the binary op tables). +const ( + unarySiLU = 6 + unaryGELU = 9 + unaryReLU2 = 10 // squared ReLU (Nemotron); cuda-only + unarySigmoid = 11 // plain sigmoid (Qwen2-MoE shared-expert gate); cuda-only + unarySoftplus = 12 // softplus (Mamba Δ); cuda-only + unaryReLU = 13 // plain ReLU max(x,0) (T5 v1.0 FFN); cuda-only + binaryAdd = 0 + binaryMul = 2 + binarySwiGLU = 6 // fused silu(a)·b — one dispatch instead of SiLU+Mul (§T613) +) + +// buffer is device-resident storage (metal/cuda/vulkan DeviceBuffer). Decoder arithmetic buffers +// are f32. The opt-in Metal f16-KV capability uses this same lifetime-only interface for K/V cache +// storage; it never calls UploadF32/DownloadF32 on those half buffers. +type buffer interface { + UploadF32([]float32) error + DownloadF32([]float32) error + Release() +} + +// qweight is a backend-resident quantized weight (metal/vulkan *ResidentQWeight) — opaque to the +// core except for Close. +type qweight interface { + Close() error +} + +// binaryNRecorder is an OPTIONAL recorder capability (same shape as quantAccRecorder): an +// elementwise op bounded to the first n elements. Decoder buffers are sized for the whole context, +// so an unbounded add over one decoded row touches ctx x width elements instead of width. Recorders +// that do not implement it keep the unbounded call and stay correct, just slower. +type binaryNRecorder interface { + BinaryN(a, b, o buffer, op, n int) error +} + +// f16KVRecorder is the optional narrow capability needed by an f16 KV cache: convert f32 projection +// rows into retained binary16 cache storage, then read that cache in attention while Q/O and all +// accumulators remain f32. Keeping it outside recorder preserves every existing backend and default +// decoder path; only NewQuantF16KV supplies both this capability and an f16 cache allocator. +type f16KVRecorder interface { + Copy2DF32ToF16Pair( + kSrc buffer, kSrcOff, kSrcStride int, + vSrc buffer, vSrcOff, vSrcStride int, + kDst buffer, kDstOff, kDstStride int, + vDst buffer, vDstOff, vDstStride int, + rows, rowFloats int, + ) error + MHAF16KV(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error +} + +// ropeF16KVAppendRecorder is the optional Metal single-token capability that rotates contiguous +// Q/K projection rows while appending K/V directly to their binary16 cache row. Attention remains +// a separate dispatch because Metal provides no grid-wide barrier between the append and MHA work. +type ropeF16KVAppendRecorder interface { + RoPEF16KVAppend(q, k, v, inv, kCache, vCache buffer, + headsQ, headsK, hd, half, pos, cacheOff int, posDiv float32, + ) error +} + +// ropePairSplitRecorder is the Metal-only fused-QKV post-projection capability. It preserves the +// ordinary recorder interface for every other backend while allowing a grouped projection to +// rotate q/k and materialize contiguous q/k/v in one dispatch instead of rotate plus three copies. +type ropePairSplitRecorder interface { + RoPEPairSplit(qkv, inv, q, k, v buffer, + seq, stride, headsQ, offQ, headsK, offK, hd, half, vOff, vDim, pos int, posDiv float32, + ) error +} + +// ropePairF16KVAppendRecorder is the grouped-QKV counterpart of ropeF16KVAppendRecorder. It keeps +// Q/K in the combined projection row while appending its K/V bands directly into half caches. +type ropePairF16KVAppendRecorder interface { + RoPEPairF16KVAppend(qkv, inv, kCache, vCache buffer, + stride, headsQ, offQ, headsK, offK, hd, half, vOff, vDim, pos, cacheOff int, posDiv float32, + ) error +} + +// linear records o[m,N] = x[m,K]·W for one projection — either an f32 device weight (MatMul) or a +// resident quantized weight (QMatMulResident). This is what lets ONE Decoder core serve both the +// f32 and the quantized model (§T415). +type linear interface { + record(r recorder, x, o buffer, m int) error + // recordAdd records dst += x·W (the residual-add epilogue, §T613). The f32 weight fuses + // the add into the matmul (one dispatch); the quantized weight has no accumulate path, + // so it multiplies into scratch and adds — the pre-fusion two-dispatch shape. + recordAdd(r recorder, x, scratch, dst buffer, m, width int) error + // recordMoE is record() with per-token MoE expert gating: on a recorder that implements the gated + // path (CUDA Q8), a token whose routing weight moeGate[row*gateStride+ex]==0 skips the weight-stream + // (bit-identical to record + the RowAxpy weight-0 combine, ~E/K faster). Dense-falls-back otherwise. + recordMoE(r recorder, x, o buffer, m int, moeGate buffer, gateStride, ex int) error +} + +// moeGatedRecorder is the optional gated-expert-GEMV extension, implemented only by the CUDA recorder for +// Q8 experts; quantLinear.recordMoE type-asserts it and dense-falls-back when absent. +type moeGatedRecorder interface { + QMatMulResidentMoE(x buffer, w qweight, o buffer, m int, moeGate buffer, gateStride, ex int) error +} + +// quantAccRecorder is the optional residual-add-FUSED quant-GEMV extension: dst += x·dequant(w) in ONE +// kernel (beta=1), instead of GEMV→scratch (beta=0) plus a separate Binary add. Implemented by the CUDA +// recorder for the Q8 decode GEMV; quantLinear.recordAdd type-asserts it for m==1 and falls back to the +// two-op shape when absent or for a format without a fused path (errQuantAccUnsupported). +type quantAccRecorder interface { + QMatMulResidentAcc(x buffer, w qweight, dst buffer, m int) error +} + +// dependencyBarrierRecorder is implemented by the Metal adapter. Ordinary/profiling Metal +// recorders make Barrier a no-op; the decode-only concurrent recorder emits a buffer-scope fence. +// Other backends preserve their established command ordering without implementing this extension. +type dependencyBarrierRecorder interface { + Barrier() error +} + +func recordBarrier(r recorder) error { + if br, ok := r.(dependencyBarrierRecorder); ok { + return br.Barrier() + } + return nil +} + +// errQuantAccUnsupported signals a weight format with no fused residual-add GEMV — recordAdd then uses +// the GEMV-into-scratch + Binary-add path. +var errQuantAccUnsupported = errors.New("llamagpu: no fused quant residual-add for this weight format") + +// swiGLUHalvesRecorder is the optional column-fused gate|up SwiGLU extension: given a [rows, 2·hidden] +// gate|up GEMV output, compute out[r,i] = silu(gu[r,i])·gu[r,hidden+i] in ONE op — the launch-count +// twin of the residual fusion for the FFN. Implemented by the CUDA recorder; recordFFN uses it when a +// block carries the fused wGU weight (built only when ops.fusedGateUp), else the two-GEMV path. +type swiGLUHalvesRecorder interface { + SwiGLUHalves(gu, out buffer, rows, hidden int) error +} + +// geGLUHalvesRecorder is the GeGLU (Gemma) twin: out[r,i] = GELU(gu[r,i])·gu[r,hidden+i] over a +// column-fused gate|up buffer, one op replacing two GEMVs + Unary GELU + Binary mul. +type geGLUHalvesRecorder interface { + GeGLUHalves(gu, out buffer, rows, hidden int) error +} + +type f32Linear struct { + w buffer + k, n int +} + +func (l f32Linear) record(r recorder, x, o buffer, m int) error { + return r.MatMul(x, l.w, o, m, l.k, l.n) +} + +func (l f32Linear) recordAdd(r recorder, x, _, dst buffer, m, _ int) error { + return r.MatMulAcc(x, l.w, dst, m, l.k, l.n) +} + +type quantLinear struct{ w qweight } + +func (l quantLinear) record(r recorder, x, o buffer, m int) error { + return r.QMatMulResident(x, l.w, o, m) +} + +func (l quantLinear) recordAdd(r recorder, x, scratch, dst buffer, m, width int) error { + // Decode / small batch (m<6): fuse the residual add into the quant kernel (beta=1) — ONE launch into + // dst instead of QMatMulResident(beta=0)→scratch + a separate Binary add, saving the add launch and + // the scratch HBM round-trip. The old code fused only m==1 because a beta=1 GEMV re-streams the weight + // M× per output (undoing the weight-read-once routing); that objection is void now that the Acc paths + // route small batches to the SAME weight-read-once MT kernel their plain twins use (Q8 #1011/#1012, + // K-quants #1013), so beta=1 does strictly less work than beta=0+add. Gate m<6: below it every quant + // uses the MT/GEMV for BOTH plain and Acc, so fused is guaranteed ≤ split. At m>=6 the plain path may + // switch to a faster kernel the Acc path lacks (Q8→int8-MMQ at m>=6, K-quants→WMMA at m>=48), so the + // split — which gets that faster kernel for the GEMM — is kept. Bit-identical either way (the Acc MT is + // bit-identical to plain MT, and beta=1 folds the add exactly; TestCUDAKQuantAccMTMatchesGEMV @m=4). + if m < 6 { + if acc, ok := r.(quantAccRecorder); ok { + if err := acc.QMatMulResidentAcc(x, l.w, dst, m); err != errQuantAccUnsupported { + return err + } + } + } + // Bound the add to the m rows actually written. Without this it runs over the whole + // context-sized buffer — 1024x the traffic at ctx=1024, and it dominated decode. + if bn, ok := r.(binaryNRecorder); ok && width > 0 { + e := r.QMatMulResident(x, l.w, scratch, m) + e = firstErr(e, recordBarrier(r)) + return firstErr(e, bn.BinaryN(dst, scratch, dst, binaryAdd, m*width)) + } + e := r.QMatMulResident(x, l.w, scratch, m) + e = firstErr(e, recordBarrier(r)) + return firstErr(e, r.Binary(dst, scratch, dst, binaryAdd)) +} + +func (l f32Linear) recordMoE(r recorder, x, o buffer, m int, _ buffer, _, _ int) error { + return l.record(r, x, o, m) // f32 MoE experts (rare) use the dense path +} + +func (l quantLinear) recordMoE(r recorder, x, o buffer, m int, moeGate buffer, gateStride, ex int) error { + if mg, ok := r.(moeGatedRecorder); ok { + return mg.QMatMulResidentMoE(x, l.w, o, m, moeGate, gateStride, ex) + } + return l.record(r, x, o, m) +} + +// recorder is one open batched command buffer (metal.Recorder / vulkan.Recorder); the adapter +// asserts the buffer args back to its concrete type. +type recorder interface { + RMSNorm(x, g, o buffer, rows, dim int, eps float32) error + LayerNorm(x, g, b, o buffer, rows, dim int, eps float32) error + AddBias(x, b, o buffer, rows, n int) error + MatMul(a, b, c buffer, m, k, n int) error + // MatMulAcc records c += a·b — the residual-add epilogue (§T613): the projection lands + // directly in the running residual stream, saving the separate elementwise-add dispatch. + MatMulAcc(a, b, c buffer, m, k, n int) error + RoPE(q, inv, o buffer, seq, width, heads, hd, half, pos int, posDiv float32) error + // RoPEAt rotates a sub-row view living at float-element offset `off` inside a wider + // buffer — the q/k bands of a fused QKV projection output (§T613). width acts as the + // ROW STRIDE; only the first heads·hd columns of each row are rotated, in place. + RoPEAt(q, inv, o buffer, off, seq, width, heads, hd, half, pos int, posDiv float32) error + // RoPEPair rotates the q band (headsQ heads at offQ) AND the k band (headsK at offK) of + // a fused QKV buffer in ONE dispatch (§T613), rows `stride` floats wide, in place. + RoPEPair(qkv, inv buffer, seq, stride, headsQ, offQ, headsK, offK, hd, half, pos int, posDiv float32) error + // RoPEPartialPair is the partial-rotary sibling of RoPEPair: only the first rotaryDim + // channels of each head in the q and k bands rotate (GPT-NeoX/Phi/StableLM partial + // rotary). inv must be the [rotaryDim/2] frequency table. Implemented on cuda; metal and + // vulkan return an unsupported error until a partial-rotary decoder targets them. + RoPEPartialPair(qkv, inv buffer, seq, stride, headsQ, offQ, headsK, offK, hd, rotaryDim, pos int, posDiv float32) error + Blit(src buffer, srcOff int, dst buffer, dstOff, n int) error + // Copy2D moves a strided rows×rowFloats sub-matrix (the fused-QKV band extraction, §T613): + // row r copies from src[srcOff+r·srcStride:] to dst[dstOff+r·dstStride:]. + Copy2D(src buffer, srcOff, srcStride int, dst buffer, dstOff, dstStride, rows, rowFloats int) error + MHA(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error + // MHACap is MHA with a Gemma-2 attention-logit soft-cap (cap·tanh(score·scale/cap) before the + // mask+softmax). Implemented on cuda; metal/vulkan return unsupported until a softcap decoder + // targets them (the RoPEPartialPair pattern). + MHACap(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale, cap float32) error + // MHAALiBi is MHA with an ALiBi position bias (MPT): slopeₕ·(j−qabs) added to each scaled score + // before the mask+softmax; slopes holds `heads` per-head slopes. cuda-only (metal/vulkan stub). + MHAALiBi(q, k, v, o, slopes buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error + // MHABias is MHA with a PRECOMPUTED per-head additive bias [heads,sq,sk] added to the scaled scores + // before softmax — the T5 relative-position bias (a learned table, not ALiBi's slope). No RoPE; + // causal==0 unmasks all keys (T5's bidirectional encoder). cuda-only (metal/vulkan stub). + MHABias(q, k, v, o, bias buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error + // MoEGate writes Mixtral top-k renormalized routing weights [rows·e] from router logits [rows·e]. + // RowAxpy accumulates dst[r,:] += arow[r]·src[r,:] (the MoE weighted combine). Both cuda-only. + MoEGate(logits, weights buffer, rows, e, k, raw int, scale float32) error + RowAxpy(dst, src, arow buffer, rows, cols int) error + // MHARect is multi-head attention with a rectangular head shape: query/key share head width dqk, + // value has a different head width dv (o is heads·dv wide). The DeepSeek-V2 MLA shape. cuda-only. + MHARect(q, k, v, o buffer, sq, sk, heads, kvHeads, dqk, dv, causal, window int, scale float32) error + // SSMStep advances ONE Mamba selective-scan timestep (decode): given u/Δ/A/B/C[d,·] and the skip + // D[d], it updates the [d,n] recurrent state h IN PLACE and writes y[d]. Conv1DStep advances one + // causal depthwise-conv step: out[c]=Σ_k w[c,k]·window[c,k]+b[c], then shifts the [d,k-1] state and + // appends x. Both carry per-block state across calls (no KV cache) — the GPU Mamba block. cuda-only. + SSMStep(u, delta, a, b, c, dskip, h, y buffer, d, n int) error + Conv1DStep(x, w, b, state, out buffer, d, k int) error + // SSDStep advances ONE Mamba-2 SSD timestep: H heads with a SCALAR per-head decay a=exp(Δ[h]·A[h]) + // and B/C shared across a group (g=h/(H/G)); updates the per-head state[H,N,P] in place and writes + // y[H·P]. x/y are [H·headDim]=intermediate; delta/a/dskip are [H]; b/c are [G·N]. cuda-only. + SSDStep(x, delta, a, b, c, dskip, state, y buffer, heads, headDim, groups, n int) error + // WKVStep advances ONE RWKV-4 WKV recurrence timestep: k/v/w/u/out are [d] (w = per-channel decay + // exp(WLog), u = current-token bonus); the running state aa/bb/pp [d] is updated in place (fresh + // sequence: aa=bb=0, pp=−1e38). No KV cache — RWKV's O(1) recurrent inference. cuda-only. + WKVStep(k, v, w, u, aa, bb, pp, out buffer, d int) error + Unary(x, o buffer, op int) error + Binary(a, b, o buffer, op int) error + QMatMulResident(x buffer, w qweight, o buffer, m int) error + // Commit submits without waiting; Wait blocks until done (§T614 encode-overlap split). + // Backends without async submit implement Commit as a deferral and Wait as Finish. + Commit() error + Wait() error + Finish() error + Free() +} + +// backendOps is the per-backend constructor vtable the adapters fill in. +type backendOps struct { + name string + newBuffer func([]float32) (buffer, error) + newF16KVBuffer func(int) (buffer, error) // nil unless this constructor opts into binary16 K/V storage + newRecorder func() (recorder, error) + // newDecodeRecorder optionally specializes only the single-token Step graph. Prefill and every + // other recorder consumer retain newRecorder; nil falls back to newRecorder. + newDecodeRecorder func() (recorder, error) + uploadQWeight func(weight []byte, qt uint32, n, k int) (qweight, error) // resident quantized upload + groupQWeights func(weights []qweight) (qweight, error) // optional heterogeneous prefill group + // quantizeF32 quantizes a host f32 weight [in,out] to a resident Q8_0 device weight in one call + // (no pre-quantized ggml bytes needed). Set only by the cuda Q8 recurrent-decode constructors — + // nil elsewhere, so the caller keeps the f32 path. Lets a decoder trade the f32 weight's bandwidth + // (the bottleneck for the weight-bound recurrent decode) for a ~4× smaller Q8 GEMV. + quantizeF32 func(w *tensor.Tensor) (qweight, error) + // asyncEncode: the backend can encode a second recorder while one executes (§T614). + // Metal command buffers are independent objects → true; the vulkan bridge keeps ONE + // global recording context → false (pre-encoding there would clobber the in-flight one). + asyncEncode bool + // fusedGateUp: the backend's recorder implements swiGLUHalvesRecorder, so a SwiGLU FFN's + // ffn_gate|ffn_up can be column-fused into ONE GEMV (into the gu buffer) + one SwiGLUHalves op, + // instead of two GEMVs + a Binary SwiGLU. Set by the CUDA SwiGLU constructors; nil elsewhere. + fusedGateUp bool + // fusedF32QKV opts a backend into GPT's single-resident-weight QKV path. The recorder must + // implement f32QKVBandsRecorder so large prefills can address each output band without copies. + fusedF32QKV bool + // fusedBiasGELU opts GPT into the recorder's bounded bias-plus-exact-GELU epilogue. Recorders + // without the optional capability retain the portable AddBias plus Unary chain. + fusedBiasGELU bool + // eagerFullLogits retains the historical Ctx×Vocab output scratch for same-binary benchmarks. + // Production constructors leave it false and lazily materialize multi-row StepN output instead. + eagerFullLogits bool +} + +// moeFFN is one sparse-MoE expert: a SwiGLU FFN (gate/up/down) with its own weights. +type moeFFN struct { + wG, wU, wD linear + wGU linear // column-fused ffn_gate|ffn_up (built only when ops.fusedGateUp); nil ⇒ wG/wU path +} + +type block struct { + wq, wk, wv, wo, wG, wU, wD linear + // wqkv is the fused QKV projection (§T613): one [in, D+2·kvDim] weight whose output row + // is [q | k | v]. Non-nil = Step/StepN run the fused chain (one matmul instead of three); + // nil = fall back to wq/wk/wv (quantized blocks whose projections mix quant types). + wqkv linear + // wqkvPrefill is the Metal-only mixed-quant companion: StepN uses one exact post-expansion + // projection, while Step keeps wq/wk/wv and therefore preserves the established decode kernels. + wqkvPrefill linear + // wGU is the column-fused ffn_gate|ffn_up projection: one [in, 2·hidden] weight whose output row + // is [gate | up]. Non-nil (built only when ops.fusedGateUp, i.e. CUDA SwiGLU) = recordFFN runs one + // GEMV into the gu buffer + one SwiGLUHalves; nil = the separate wG/wU two-GEMV SwiGLU path. + wGU linear + gAttn, gFFN, kC, vC buffer + bAttn, bFFN buffer // LayerNorm betas (nil ⇒ RMSNorm, the Llama default) + // projection biases (nil ⇒ no bias): the fused-QKV bias [D+2·kvDim], the o-proj bias [D], + // and the GELU-MLP fc [hidden] / proj [D] biases — GPT-NeoX/Phi/StarCoder2 carry them. + qkvBias, oBias, fcBias, projBias buffer + qN, kN buffer // per-head query/key RMSNorm gains [dk] (Qwen3); nil ⇒ no QK-norm + gAttn2, gFFN2 buffer // sandwich-norm OUTPUT gains (Gemma2 post_attention/post_feedforward); nil otherwise + moeRouter linear // sparse-MoE gate: dim → nExperts logits (Mixtral); nil unless d.moe + moeExperts []moeFFN // the per-expert SwiGLU FFNs (Mixtral); len == d.nExperts + moeShared moeFFN // Qwen2-MoE shared expert (SwiGLU on every token); zero-value when absent + moeSharedGate linear // Qwen2-MoE shared-expert gate: dim → 1 logit, sigmoid-scaled; nil when absent + wqA, wqB, wkvA, wkvB linear // DeepSeek-V2 MLA low-rank projections (q_a/q_b/kv_a/kv_b); nil unless d.mla + gQA, gKvA buffer // MLA latent RMSNorm gains (q_a_layernorm / kv_a_layernorm) + mInX, mInZ, mDtLow, mDtProj, mBProj, mCProj, mOutProj linear // Mamba block projections; nil unless d.mamba + mConvW, mConvB, mA, mDskip, mDtBias buffer // Mamba conv filters+bias, A=−exp(A_log) [dInner,N], D skip [dInner], Δ bias [dInner] + mConvState, mSsmState buffer // per-block Mamba decode state: conv window [dInner,dConv-1], ssm state [dInner,N] + mDtNorm, mBNorm, mCNorm buffer // Jamba mixer RMSNorm gains on Δ_low [dtRank] / B [N] / C [N]; nil for plain Mamba (recordMambaBlock applies them only when set) + jMamba bool // Jamba: this layer's mixer is Mamba (else NoPE attention); only meaningful when d.jamba + m2InProj, m2OutProj linear // Mamba-2 fused in_proj [d,projSize] / out_proj [inter,d]; nil unless d.mamba2 + m2ConvW, m2ConvB, m2A, m2DtBias, m2Dskip, m2NormW buffer // Mamba-2 conv filters+bias [convDim,·], A/Δbias/D [numHeads], gated-RMSNorm gain [inter] + m2ConvState, m2SsdState buffer // Mamba-2 state: conv window [convDim,dConv-1], SSD state [numHeads·N·headDim] + // RWKV-4 (nil unless d.rwkv): time-mix reuses wq/wk/wv/wo = Wr/Wk/Wv/Wo, channel-mix reuses + // wG/wD = CWk/CWv; LN1 = gAttn/bAttn, LN2 = gFFN/bFFN. rwCWr is the channel-mix receptance. + rwCWr linear // channel-mix receptance CWr [Dim,Dim] + rwMuR, rwMuK, rwMuV, rwCMuR, rwCMuK buffer // token-shift interpolators μ [Dim] + rwOmR, rwOmK, rwOmV, rwOmCR, rwOmCK buffer // (1−μ) precomputed [Dim] + rwW, rwU buffer // WKV per-channel decay w=exp(WLog) [Dim], bonus u [Dim] + rwPrevTM, rwPrevCM, rwAA, rwBB, rwPP buffer // per-block RWKV recurrent state [Dim] (token-shift + WKV aa/bb/pp) +} + +// Decoder holds a Llama's weights + KV cache as device-resident buffers and runs one batched decode +// step per token. Not safe for concurrent use — one Decoder per goroutine. Release when done. +type Decoder struct { + ops backendOps + d, h, kvH, dk, kvDim, half, hidden, v, maxLen int + qDim int // attention Q/O width = h·dk; == d (model dim) unless head_dim is decoupled (Gemma) + rotaryDim int // rotated channels/head; == dk for full RoPE (Llama), < dk for partial rotary (GPT-NeoX/Phi/StableLM) + lnBias bool // true ⇒ LayerNorm-with-bias norms (StableLM/Phi/StarCoder2); false ⇒ RMSNorm (Llama) + ffnGELU bool // true ⇒ 2-layer GELU MLP (fc→GELU→proj: GPT-NeoX/Phi/StarCoder2); false ⇒ SwiGLU (Llama/StableLM) + ffnReLU2 bool // true ⇒ 2-layer squared-ReLU MLP (fc→relu²→proj: Nemotron); mutually exclusive with ffnGELU + ffnGEGLU bool // true ⇒ GeGLU MLP: down(GELU(gate)⊙up) — Gemma's GELU-gated variant of SwiGLU + parallelRes bool // true ⇒ parallel residual: attn+FFN both read norm(x0), sum onto x (GPT-NeoX/Phi/Cohere); false ⇒ sequential (Llama) + parallelTwoNorm bool // parallel residual with SEPARATE attn/FFN norms both over x0 (GPT-NeoX); needs the FFN norm computed pre-attn-add + qkNorm bool // true ⇒ RMSNorm on Q and K before RoPE (Qwen3 per-head, OLMo2 full-width); false otherwise + qkNormFull bool // with qkNorm: true ⇒ one RMSNorm over the whole q/k projection (OLMo2), false ⇒ per-head (Qwen3) + postNorm bool // true ⇒ post-norm blocks (OLMo2): sublayer reads the raw residual, its OUTPUT is normed before the add + sandwich bool // true ⇒ sandwich norms (Gemma2): pre-norm the input (gAttn/gFFN) AND post-norm the output (gAttn2/gFFN2) + attnCap float32 // >0 ⇒ Gemma2 attention-logit soft-cap magnitude (via MHACap); 0 ⇒ plain MHA + noRope bool // true ⇒ no rotary position embedding (MPT: position enters via ALiBi only) + aliBiSlopes buffer // non-nil ⇒ per-head ALiBi slopes [heads]; attention runs via MHAALiBi (MPT) + moe bool // true ⇒ the FFN sublayer is a sparse Mixture-of-Experts (Mixtral): route → experts → combine + nExperts, topK int // MoE expert count and experts-per-token (when moe) + moeRaw bool // true ⇒ DeepSeek-V2 raw-softmax·routedScale gating (no top-k renormalization); false ⇒ Mixtral renorm + routedScale float32 // DeepSeek-V2 routed-expert weight multiplier (routed_scaling_factor); 1 otherwise + mla bool // true ⇒ DeepSeek-V2 Multi-head Latent Attention (rectangular, low-rank latent KV, decoupled RoPE) + mamba bool // true ⇒ Mamba selective-state-space blocks (no attention, no KV cache; linear-time recurrent decode) + dInner, mambaN, dConv, dtRank int // Mamba block dims: inner width (e·d), SSM state size, conv kernel, Δ low-rank + mamba2 bool // true ⇒ Mamba-2 SSD blocks (scalar-per-head decay, grouped B/C, fused in_proj, gated RMSNorm) + m2H, m2P, m2G, m2N, m2Conv, m2Inter, m2CD int // Mamba-2 dims: num_heads, head_dim, n_groups, state_size, conv kernel, intermediate, conv_dim + rwkv bool // true ⇒ RWKV-4 blocks (WKV time-mix + gated squared-ReLU channel-mix; no attention/KV cache; O(1) recurrent) + rwHidden int // RWKV channel-mix hidden width (CWk: Dim→Hidden, CWv: Hidden→Dim) + jamba bool // true ⇒ Jamba hybrid stack: each layer is Mamba-mixer (jMamba) OR NoPE GQA attention, then a per-block MoE/dense FFN; sequential like Mamba + qkNope, qkRope, vHead, qkHead int // MLA per-head dims: nope/rope query-key parts, value width, qkHead=qkNope+qkRope + qLoRA, kvLoRA int // MLA query/kv latent compression ranks (q_lora_rank / kv_lora_rank) + mlaScale, mlaPosDiv float32 // MLA pre-softmax scale (1/√qkHead default) and the QKRope rope posDiv + eps, posDiv, scale float32 + embMult float32 // gathered embeddings ×= embMult (IBM Granite EmbeddingMult); 1 for everything else + f16KV bool // opt-in: K/V cache storage is IEEE binary16; compute remains f32 + + blocks []block + out linear + gFinal, bFinal, outBias, dinv *bufSlot + dx, xn, xn2, q, k, v_, attn, ao, gate, up, mo, logits *bufSlot + logitsRows int + fullLogits growBuffer + gu *bufSlot // fused gate|up GEMV output [·, 2·hidden] (nil unless ops.fusedGateUp) + moeGate, moeW, moeCol *bufSlot // sparse-MoE scratch: router logits [·,E], routing weights [·,E], one weight column [·] + mlaCQ, mlaQ, mlaComp, mlaLatent, mlaKV, mlaAttn, mlaInv *bufSlot // MLA scratch: compressed query, fused Q, kv_a out, normed kv latent, kv_b out, attn out, QKRope freqs + mbXin, mbZ, mbXc, mbDtLow, mbDelta, mbB, mbC, mbY *bufSlot // Mamba scratch (rows=1): in_x, gate branch, conv+SiLU out, Δ_low, Δ, B, C, scan out + m2Proj, m2Z, m2XBC, m2Dt, m2Xc, m2X, m2B, m2C, m2Y *bufSlot // Mamba-2 scratch (rows=1): in_proj out, gate z, conv input xBC, Δ pre-act, conv+SiLU out, value x, B, C, SSD out + rwPreLNg, rwPreLNb *bufSlot // RWKV ln0 (pre-block-0 LayerNorm) γ/β + rwMix, rwR, rwK, rwV, rwWkv, rwO, rwT1, rwT2, rwCr, rwCk, rwCv *bufSlot // RWKV scratch (rows=1): token-shift mix, r/k/v, WKV out, time-mix out, 2 mix temps, chan-mix cr, ck [Hidden], cv + qkv *bufSlot // fused QKV output rows [·, d+2·kvDim] (§T613) + pending recorder // pre-encoded next-step command buffer (§T614) + pendingPos int // position pending encodes; -1 = none + table *tensor.Tensor // token embedding, host-side gather source + invHost []float32 // RoPE inverse freqs (uploaded by allocScratch) + all []buffer + qweights []qweight // resident quantized weights (quant decoder) +} + +// bufSlot wraps a buffer so the struct fields read naturally while sharing the release list. +type bufSlot struct{ b buffer } + +// growBuffer owns one backend buffer whose capacity grows only when a caller requests more +// elements. It is deliberately kept outside Decoder.all: growth releases the old buffer before +// allocating its exact replacement, while decoder Release handles the final owner explicitly. +type growBuffer struct { + b buffer + n int +} + +func (g *growBuffer) ensure(newBuffer func([]float32) (buffer, error), n int) (buffer, error) { + if g.b != nil && g.n >= n { + return g.b, nil + } + g.release() + b, err := newBuffer(make([]float32, n)) + if err != nil { + return nil, err + } + g.b, g.n = b, n + return b, nil +} + +func (g *growBuffer) release() { + if g.b != nil { + g.b.Release() + } + g.b, g.n = nil, 0 +} + +func logitsForRows(ops backendOps, resident buffer, residentRows int, overflow *growBuffer, rows, vocab int) (buffer, error) { + if rows <= residentRows { + return resident, nil + } + return overflow.ensure(ops.newBuffer, rows*vocab) +} + +// newDecoderCommon sets the geometry, RoPE frequencies and scratch buffers shared by the f32 and +// quantized constructors. +func newDecoderCommon(cfg nlp.LlamaConfig, tokEmb *tensor.Tensor, ops backendOps) (*Decoder, error) { + d := &Decoder{ + ops: ops, + d: cfg.Dim, h: cfg.Heads, kvH: cfg.KVHeads, hidden: cfg.Hidden, v: cfg.Vocab, + maxLen: cfg.Ctx, eps: float32(cfg.Eps), table: tokEmb, + pendingPos: -1, embMult: 1, + } + if d.kvH <= 0 { + d.kvH = d.h + } + if d.h <= 0 || d.d%d.h != 0 { + return nil, fmt.Errorf("llamagpu(%s): dim %d not divisible by heads %d", ops.name, d.d, d.h) + } + d.dk = d.d / d.h + d.qDim = d.h * d.dk // = d.d for standard models; decoupled-head_dim constructors (Gemma) override + d.kvDim = d.kvH * d.dk + d.half = d.dk / 2 + d.rotaryDim = d.dk // full RoPE by default; partial-rotary constructors override before allocScratch + d.scale = float32(1.0 / math.Sqrt(float64(d.dk))) + + invF64, posDiv64 := backend.RoPEFreqs(d.dk, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, d.half) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + return d, nil +} + +// ropeQK records RoPE over the q and k bands of the fused QKV buffer in one pass: full rotary +// (RoPEPair) when rotaryDim == dk (Llama), else partial rotary (RoPEPartialPair — only the first +// rotaryDim channels of each head rotate; GPT-NeoX/Phi/StableLM). inv (d.dinv) is sized to +// rotaryDim/2 so the same buffer feeds both. +// recordFFN records the feed-forward sublayer with the residual-add epilogue (dx += ffn(xn2)): +// a 2-layer GELU MLP (fc → GELU → proj, GPT-NeoX/Phi/StarCoder2) when ffnGELU, else the SwiGLU +// (silu(gate)·up → down, Llama/StableLM). wG/wU/wD are c_fc/·/c_proj or gate/up/down accordingly. +// recordQKVProj records the fused QKV projection (xn·Wqkv) plus its optional bias — the bias is +// added BEFORE RoPE, per the [q|k|v] band layout (GPT-NeoX/Phi/StarCoder2 have biased q/k/v). +// recordAttnNorm produces the attention sublayer's input in d.xn. Pre-norm (Llama & most): xn = +// norm(x0) (plus norm2(x0)→xn2 for two-norm parallel). Post-norm (OLMo2): attention reads the RAW +// residual, so xn is just a copy of x0 and the norm happens later on the attention output. +func (d *Decoder) recordAttnNorm(r recorder, b block, rows int) error { + if d.postNorm { + return r.Blit(d.dx.b, 0, d.xn.b, 0, rows*d.d) + } + e := d.norm(r, d.dx.b, b.gAttn, b.bAttn, d.xn.b, rows) + if d.parallelTwoNorm { // norm2(x0) BEFORE the attn add, for the FFN branch + e = firstErr(e, d.norm(r, d.dx.b, b.gFFN, b.bFFN, d.xn2.b, rows)) + } + return e +} + +func (d *Decoder) recordQKVProj(r recorder, b block, rows int) error { + return d.recordQKVProjWith(r, b.wqkv, b.qkvBias, rows) +} + +func (d *Decoder) recordQKVProjWith(r recorder, w linear, bias buffer, rows int) error { + e := w.record(r, d.xn.b, d.qkv.b, rows) + if bias != nil { + e = firstErr(e, r.AddBias(d.qkv.b, bias, d.qkv.b, rows, d.qDim+2*d.kvDim)) + } + return e +} + +// recordQKNorm records Qwen3's per-head RMSNorm on the Q and K bands of the fused QKV buffer, +// applied BEFORE RoPE. Each head's dk channels are normed independently with a shared gain — the +// reference reshapes [seq, heads·dk] → [seq·heads, dk], RMSNorms, reshapes back. The heads of one +// token are contiguous but tokens are stride-separated inside qkv, so we Copy2D the band out to a +// packed [seq, band] scratch (reusing d.q / d.k, unused in the fused path until after RoPE), +// RMSNorm it as seq·heads rows of dk, and Copy2D it back. No-op when qkNorm is unset. +func (d *Decoder) recordQKNorm(r recorder, b block, seq, stride int) error { + if !d.qkNorm { + return nil + } + // rows×dim for the RMSNorm: OLMo2 norms the WHOLE q/k projection as one vector (rows=seq, + // dim=qDim/kvDim, gain [qDim]/[kvDim]); Qwen3 norms each head (rows=seq·heads, dim=dk, gain [dk]). + qRows, qCols, kRows, kCols := seq*d.h, d.dk, seq*d.kvH, d.dk + if d.qkNormFull { + qRows, qCols, kRows, kCols = seq, d.qDim, seq, d.kvDim + } + e := r.Copy2D(d.qkv.b, 0, stride, d.q.b, 0, d.qDim, seq, d.qDim) // extract Q band [seq, qDim] + e = firstErr(e, r.RMSNorm(d.q.b, b.qN, d.q.b, qRows, qCols, d.eps)) + e = firstErr(e, r.Copy2D(d.q.b, 0, d.qDim, d.qkv.b, 0, stride, seq, d.qDim)) // write Q band back + e = firstErr(e, r.Copy2D(d.qkv.b, d.qDim, stride, d.k.b, 0, d.kvDim, seq, d.kvDim)) + e = firstErr(e, r.RMSNorm(d.k.b, b.kN, d.k.b, kRows, kCols, d.eps)) + e = firstErr(e, r.Copy2D(d.k.b, 0, d.kvDim, d.qkv.b, d.qDim, stride, seq, d.kvDim)) + return e +} + +// recordOProj records the attention output projection with its residual-add epilogue (dx += +// attn·Wo) plus the optional o-bias broadcast onto the residual stream. +// recordMHA runs attention against the Decoder's configured cache storage. The f16 capability is +// selected only by the opt-in constructor; all established constructors still take the exact f32 +// calls below. Initial batched prefill can explicitly pass f16=false to recordMHAStorage so FlashMM +// consumes its already-resident f32 K/V rows while those rows are also converted into the cache. +func (d *Decoder) recordMHA(r recorder, q, kC, vC, o buffer, sq, sk int) error { + return d.recordMHAStorage(r, q, kC, vC, o, sq, sk, d.f16KV) +} + +func (d *Decoder) recordMHAStorage(r recorder, q, kC, vC, o buffer, sq, sk int, f16 bool) error { + if d.aliBiSlopes != nil { + return r.MHAALiBi(q, kC, vC, o, d.aliBiSlopes, sq, sk, d.qDim, d.h, d.kvH, d.dk, 1, 0, d.scale) + } + if d.attnCap > 0 { + return r.MHACap(q, kC, vC, o, sq, sk, d.qDim, d.h, d.kvH, d.dk, 1, 0, d.scale, d.attnCap) + } + if f16 { + fr, ok := r.(f16KVRecorder) + if !ok { + return fmt.Errorf("llamagpu(%s): f16 KV cache requires recorder conversion/attention capability", d.ops.name) + } + return fr.MHAF16KV(q, kC, vC, o, sq, sk, d.qDim, d.h, d.kvH, d.dk, 1, 0, d.scale) + } + return r.MHA(q, kC, vC, o, sq, sk, d.qDim, d.h, d.kvH, d.dk, 1, 0, d.scale) +} + +func (d *Decoder) recordKVPairBlit(r recorder, + kSrc buffer, kSrcOff int, vSrc buffer, vSrcOff int, + kDst, vDst buffer, dstOff, n int, +) error { + if !d.f16KV { + return firstErr( + r.Blit(kSrc, kSrcOff, kDst, dstOff, n), + r.Blit(vSrc, vSrcOff, vDst, dstOff, n), + ) + } + fr, ok := r.(f16KVRecorder) + if !ok { + return fmt.Errorf("llamagpu(%s): f16 KV cache requires recorder conversion capability", d.ops.name) + } + return fr.Copy2DF32ToF16Pair( + kSrc, kSrcOff, n, + vSrc, vSrcOff, n, + kDst, dstOff, n, + vDst, dstOff, n, + 1, n, + ) +} + +func (d *Decoder) recordKVPairCopy2D(r recorder, + kSrc buffer, kSrcOff, kSrcStride int, + vSrc buffer, vSrcOff, vSrcStride int, + kDst, vDst buffer, dstOff, dstStride, rows, rowFloats int, +) error { + if !d.f16KV { + return firstErr( + r.Copy2D(kSrc, kSrcOff, kSrcStride, kDst, dstOff, dstStride, rows, rowFloats), + r.Copy2D(vSrc, vSrcOff, vSrcStride, vDst, dstOff, dstStride, rows, rowFloats), + ) + } + fr, ok := r.(f16KVRecorder) + if !ok { + return fmt.Errorf("llamagpu(%s): f16 KV cache requires recorder conversion capability", d.ops.name) + } + return fr.Copy2DF32ToF16Pair( + kSrc, kSrcOff, kSrcStride, + vSrc, vSrcOff, vSrcStride, + kDst, dstOff, dstStride, + vDst, dstOff, dstStride, + rows, rowFloats, + ) +} + +func (d *Decoder) recordOProj(r recorder, b block, rows int) error { + if d.postNorm || d.sandwich { + // Output-normed residual: ao = attn·Wo → norm(ao) → dx += ao. OLMo2 post-norm reuses gAttn; + // Gemma2 sandwich has a SEPARATE post_attention_layernorm gain (gAttn2), the input already + // having been pre-normed with gAttn. + g, bta := b.gAttn, b.bAttn + if d.sandwich { + g, bta = b.gAttn2, nil + } + return firstErr( + b.wo.record(r, d.attn.b, d.ao.b, rows), + d.norm(r, d.ao.b, g, bta, d.ao.b, rows), + d.binElem(r, d.dx.b, d.ao.b, d.dx.b, binaryAdd, rows, d.d), + ) + } + e := b.wo.recordAdd(r, d.attn.b, d.ao.b, d.dx.b, rows, d.d) + if b.oBias != nil { + e = firstErr(e, r.AddBias(d.dx.b, b.oBias, d.dx.b, rows, d.d)) + } + return e +} + +// recordDownProj records the FFN down-projection's residual epilogue: dx += act·Wdown, either fused +// (pre-norm) or, for OLMo2 post-norm, as mo = act·Wdown → post_feedforward_layernorm(mo) → dx += mo. +func (d *Decoder) recordDownProj(r recorder, b block, rows int) error { + if d.postNorm || d.sandwich { + g, bta := b.gFFN, b.bFFN + if d.sandwich { + g, bta = b.gFFN2, nil // Gemma2 post_feedforward_layernorm (input pre-normed with gFFN) + } + return firstErr( + b.wD.record(r, d.gate.b, d.mo.b, rows), + d.norm(r, d.mo.b, g, bta, d.mo.b, rows), + d.binElem(r, d.dx.b, d.mo.b, d.dx.b, binaryAdd, rows, d.d), + ) + } + return b.wD.recordAdd(r, d.gate.b, d.mo.b, d.dx.b, rows, d.d) +} + +// recordLogits records the final projection to logits (xn·Out) plus the optional output bias +// (Phi's biased untied lm_head; nil ⇒ no bias, as for Llama/StableLM/StarCoder2). +func (d *Decoder) recordLogits(r recorder, logits buffer, rows int) error { + e := d.out.record(r, d.xn.b, logits, rows) + if d.outBias != nil { + e = firstErr(e, r.AddBias(logits, d.outBias.b, logits, rows, d.v)) + } + return e +} + +// recordFFNSublayer records the FFN norm + the FFN with its residual-add epilogue. Sequential +// (Llama/StableLM/StarCoder2): norm the post-attention residual, then FFN(that). PARALLEL residual +// (Phi/Cohere one-norm): the FFN reuses the SAME normed input as attention (d.xn = norm(x0), which +// survives attention since recordOProj writes dx, not xn) and adds onto the residual in parallel — +// no second norm. Both leave dx += FFN(·). +func (d *Decoder) recordFFNSublayer(r recorder, b block, rows int) error { + if b.moeRouter != nil { + // sparse MoE — per BLOCK (DeepSeek-V2 mixes dense and MoE layers): pre-norm, then route → + // per-expert SwiGLU → weighted combine (+ optional shared expert). + if e := d.norm(r, d.dx.b, b.gFFN, b.bFFN, d.xn2.b, rows); e != nil { + return e + } + return d.recordMoE(r, b, d.xn2.b, rows) + } + if d.postNorm { + return d.recordFFN(r, b, d.dx.b, rows) // post-norm: FFN reads the raw residual; recordDownProj norms its output + } + if d.parallelRes { + return d.recordFFN(r, b, d.xn.b, rows) // parallel: FFN(norm(x0)), same norm as attn + } + if d.parallelTwoNorm { + return d.recordFFN(r, b, d.xn2.b, rows) // two-norm parallel: xn2 = norm2(x0), computed pre-attn-add + } + if e := d.norm(r, d.dx.b, b.gFFN, b.bFFN, d.xn2.b, rows); e != nil { + return e + } + if e := recordBarrier(r); e != nil { + return e + } + return d.recordFFN(r, b, d.xn2.b, rows) +} + +// recordMoE records the sparse Mixture-of-Experts FFN sublayer over the normed input `in`: +// gate logits = in·Router → top-k renormalized weights (MoEGate) → for every expert, SwiGLU(in) +// scaled by that expert's per-token weight and accumulated straight onto the residual dx (RowAxpy). +// This is DENSE eval (all experts run); the non-selected experts get weight 0, so it is exact — +// the sparse top-k gather is a later optimization. The batched command buffer stays valid because +// the routing weights are computed on-device, not by host control flow. +func (d *Decoder) recordMoE(r recorder, b block, in buffer, rows int) error { + e := b.moeRouter.record(r, in, d.moeGate.b, rows) // [rows, nExperts] router logits + e = firstErr(e, r.MoEGate(d.moeGate.b, d.moeW.b, rows, d.nExperts, d.topK, boolToInt(d.moeRaw), d.routedScale)) + sh, canFuse := r.(swiGLUHalvesRecorder) + for ex := range b.moeExperts { + xp := b.moeExperts[ex] + if xp.wGU != nil && canFuse { + // Column-fused expert gate|up: one gated GEMV into gu = [·, 2·hidden] then SwiGLU over its + // halves — halves the per-expert projection launches (the loop launches every expert, most + // gated-out but still launched). Gating zeros non-routed rows, so SwiGLUHalves(0)=0. Bit-exact. + e = firstErr(e, + xp.wGU.recordMoE(r, in, d.gu.b, rows, d.moeW.b, d.nExperts, ex), + sh.SwiGLUHalves(d.gu.b, d.gate.b, rows, d.hidden), + xp.wD.recordMoE(r, d.gate.b, d.mo.b, rows, d.moeW.b, d.nExperts, ex), + r.Copy2D(d.moeW.b, ex, d.nExperts, d.moeCol.b, 0, 1, rows, 1), + r.RowAxpy(d.dx.b, d.mo.b, d.moeCol.b, rows, d.d), + ) + continue + } + e = firstErr(e, + xp.wG.recordMoE(r, in, d.gate.b, rows, d.moeW.b, d.nExperts, ex), + xp.wU.recordMoE(r, in, d.up.b, rows, d.moeW.b, d.nExperts, ex), + d.binElem(r, d.gate.b, d.up.b, d.gate.b, binarySwiGLU, rows, d.hidden), // silu(gate)·up + xp.wD.recordMoE(r, d.gate.b, d.mo.b, rows, d.moeW.b, d.nExperts, ex), // expert output → mo + r.Copy2D(d.moeW.b, ex, d.nExperts, d.moeCol.b, 0, 1, rows, 1), // extract weight column [rows] + r.RowAxpy(d.dx.b, d.mo.b, d.moeCol.b, rows, d.d), // dx += w_ex · expert_ex + ) + } + if b.moeShared.wG != nil { + // Shared expert run on every token. Qwen2-MoE gates it by sigmoid(in·gate); DeepSeek-V2 runs + // it unconditionally (no gate), so its SwiGLU accumulates straight onto the residual. + s := b.moeShared + e = firstErr(e, + s.wG.record(r, in, d.gate.b, rows), + s.wU.record(r, in, d.up.b, rows), + d.binElem(r, d.gate.b, d.up.b, d.gate.b, binarySwiGLU, rows, d.hidden), // shared SwiGLU → gate + ) + if b.moeSharedGate != nil { // Qwen2-MoE: sigmoid-gated combine + e = firstErr(e, + s.wD.record(r, d.gate.b, d.mo.b, rows), + b.moeSharedGate.record(r, in, d.moeCol.b, rows), + r.Unary(d.moeCol.b, d.moeCol.b, unarySigmoid), + r.RowAxpy(d.dx.b, d.mo.b, d.moeCol.b, rows, d.d), // dx += sigmoid(gate)·shared_out + ) + } else { // DeepSeek-V2: ungated — dx += shared·Wdown directly + e = firstErr(e, s.wD.recordAdd(r, d.gate.b, d.mo.b, d.dx.b, rows, d.d)) + } + } + return e +} + +// recordMLAAttention records DeepSeek-V2 Multi-head Latent Attention for `rows` new tokens at +// position pos, leaving dx += attn·Wo. It fuses the reference's head-by-head build: input_layernorm → +// low-rank q/kv latent projections (each with its own RMSNorm) → decoupled RoPE on the per-head q_pe +// and the SHARED k_pe → assemble the rectangular per-head Q [qkHead] / K [qkHead] (k_pe broadcast to +// every head) and the VHead-wide V straight into the KV cache → MHARect (rectangular) → o_proj. +func (d *Decoder) recordMLAAttention(r recorder, b block, pos, rows int) error { + H, qkH, nope, rope, vH := d.h, d.qkHead, d.qkNope, d.qkRope, d.vHead + kvH := nope + vH // kv_b per-head width (k_nope + value) + qDim, vDim := H*qkH, H*vH + compW := d.kvLoRA + rope // kv_a output width: [kv_latent | k_pe] + + e := d.norm(r, d.dx.b, b.gAttn, b.bAttn, d.xn.b, rows) // input_layernorm → xn + + // Query: q = q_b_proj(q_a_layernorm(q_a_proj(xn))), then decoupled RoPE on each head's q_pe. + e = firstErr(e, + b.wqA.record(r, d.xn.b, d.mlaCQ.b, rows), + r.RMSNorm(d.mlaCQ.b, b.gQA, d.mlaCQ.b, rows, d.qLoRA, d.eps), + b.wqB.record(r, d.mlaCQ.b, d.mlaQ.b, rows), // [rows, H·qkH] + ) + for h := 0; h < H; h++ { + e = firstErr(e, r.RoPEAt(d.mlaQ.b, d.mlaInv.b, d.mlaQ.b, h*qkH+nope, rows, qDim, 1, rope, rope/2, pos, d.mlaPosDiv)) + } + + // KV: compressed = kv_a_proj(xn) → [kv_latent | k_pe]; RoPE the shared k_pe in place; norm the + // latent and expand kv = kv_b_proj(kv_a_layernorm(latent)). + e = firstErr(e, + b.wkvA.record(r, d.xn.b, d.mlaComp.b, rows), + r.RoPEAt(d.mlaComp.b, d.mlaInv.b, d.mlaComp.b, d.kvLoRA, rows, compW, 1, rope, rope/2, pos, d.mlaPosDiv), + r.Copy2D(d.mlaComp.b, 0, compW, d.mlaLatent.b, 0, d.kvLoRA, rows, d.kvLoRA), // extract kv_latent + r.RMSNorm(d.mlaLatent.b, b.gKvA, d.mlaLatent.b, rows, d.kvLoRA, d.eps), + b.wkvB.record(r, d.mlaLatent.b, d.mlaKV.b, rows), // [rows, H·kvH] + ) + + // Assemble the new tokens' K/V into the cache at row pos: K[h] = [k_nope_h | k_pe(shared)], + // V[h] = value_h. Copy2D writes rows-many rows at cache-row offset pos. + for h := 0; h < H; h++ { + e = firstErr(e, + r.Copy2D(d.mlaKV.b, h*kvH, H*kvH, b.kC, pos*qDim+h*qkH, qDim, rows, nope), // k_nope + r.Copy2D(d.mlaComp.b, d.kvLoRA, compW, b.kC, pos*qDim+h*qkH+nope, qDim, rows, rope), // shared k_pe → every head + r.Copy2D(d.mlaKV.b, h*kvH+nope, H*kvH, b.vC, pos*vDim+h*vH, vDim, rows, vH), // value + ) + } + + // Rectangular attention over the cache [0 : pos+rows], then the output projection onto dx. + return firstErr(e, + r.MHARect(d.mlaQ.b, b.kC, b.vC, d.mlaAttn.b, rows, pos+rows, H, H, qkH, vH, 1, 0, d.mlaScale), + b.wo.recordAdd(r, d.mlaAttn.b, d.mo.b, d.dx.b, rows, d.d), // dx += attn·Wo + ) +} + +// recordMambaBlock records ONE Mamba selective-state-space block for a single decode token +// (rows==1), leaving dx += mixer(RMSNorm(dx)). It mirrors nn.MambaBlock.Forward exactly, but the +// full-sequence OpConv1D/OpSSM are replaced by their per-timestep decode kernels (Conv1DStep, +// SSMStep) that carry the block's conv-window and SSM state across Step calls — no attention, no KV +// cache, O(1) work per token. The per-token recurrence is why Mamba records rows==1 only (StepN +// loops it): each token's state update depends on the previous token's. Only DtProj and the conv +// carry a bias (the HF checkpoint's other projections are bias-free); the Δ bias lands inside +// softplus. A = −exp(A_log) is precomputed on the host (b.mA). +func (d *Decoder) recordMambaBlock(r recorder, b block) error { + dI, N, K := d.dInner, d.mambaN, d.dConv + // Jamba inserts an RMSNorm on Δ_low/B/C between the x_proj split and the scan; plain Mamba passes + // nil gains, so jnorm is a no-op there (keeps this one recorder path serving both). + jnorm := func(buf, g buffer, w int) error { + if g == nil { + return nil + } + return r.RMSNorm(buf, g, buf, 1, w, d.eps) + } + // pre-norm → xn = RMSNorm(dx); the up-projection splits into the x branch and the z gate. + e := d.norm(r, d.dx.b, b.gAttn, nil, d.xn.b, 1) // RMSNorm (lnBias false for Mamba) + e = firstErr(e, + b.mInX.record(r, d.xn.b, d.mbXin.b, 1), // xin = xn·InX [dInner] + b.mInZ.record(r, d.xn.b, d.mbZ.b, 1), // z = xn·InZ [dInner] + // local mixing: one causal depthwise-conv step (state carried), then SiLU. + r.Conv1DStep(d.mbXin.b, b.mConvW, b.mConvB, b.mConvState, d.mbXc.b, dI, K), + r.Unary(d.mbXc.b, d.mbXc.b, unarySiLU), + // input-dependent Δ = softplus(dt_proj(RMSNorm?(dt_low_proj(xc))) + dt_bias). + b.mDtLow.record(r, d.mbXc.b, d.mbDtLow.b, 1), // Δ_low [dtRank] + jnorm(d.mbDtLow.b, b.mDtNorm, d.dtRank), // Jamba DtNorm + b.mDtProj.record(r, d.mbDtLow.b, d.mbDelta.b, 1), // → [dInner] + r.AddBias(d.mbDelta.b, b.mDtBias, d.mbDelta.b, 1, dI), + r.Unary(d.mbDelta.b, d.mbDelta.b, unarySoftplus), + // input-dependent SSM params B, C (Jamba RMSNorms each). + b.mBProj.record(r, d.mbXc.b, d.mbB.b, 1), // B [N] + jnorm(d.mbB.b, b.mBNorm, N), // Jamba BNorm + b.mCProj.record(r, d.mbXc.b, d.mbC.b, 1), // C [N] + jnorm(d.mbC.b, b.mCNorm, N), // Jamba CNorm + // selective scan: one recurrent timestep (state carried), y[dInner]. + r.SSMStep(d.mbXc.b, d.mbDelta.b, b.mA, d.mbB.b, d.mbC.b, b.mDskip, b.mSsmState, d.mbY.b, dI, N), + // gate y ⊙ SiLU(z), then down-project onto the residual: dx += y·OutProj. + r.Unary(d.mbZ.b, d.mbZ.b, unarySiLU), + r.Binary(d.mbY.b, d.mbZ.b, d.mbY.b, binaryMul), + b.mOutProj.recordAdd(r, d.mbY.b, d.mo.b, d.dx.b, 1, d.d), + ) + return e +} + +// resetMambaState zeroes every block's conv-window and SSM state — called at the start of a fresh +// sequence (pos==0), since the recurrence carries state across Step calls with no KV cache to reset. +func (d *Decoder) resetMambaState() error { + zc := make([]float32, d.dInner*(d.dConv-1)) + zs := make([]float32, d.dInner*d.mambaN) + for _, b := range d.blocks { + if b.mConvState == nil { // Jamba: attention layers have no conv/SSM state + continue + } + if e := b.mConvState.UploadF32(zc); e != nil { + return e + } + if e := b.mSsmState.UploadF32(zs); e != nil { + return e + } + } + return nil +} + +// recordJambaAttn records ONE Jamba attention mixer for a single decode token (rows==1), leaving +// dx += attn(RMSNorm(dx)). It is Llama-style grouped-query causal attention over the block's growing +// KV cache, but NoPE (no rotary): Jamba's Mamba layers carry position, so q/k are NEVER rotated — +// the only departure from the standard attention sublayer. Bias-free q/k/v/o. +func (d *Decoder) recordJambaAttn(r recorder, b block, pos int) error { + kvDim := d.kvDim + return firstErr( + d.recordAttnNorm(r, b, 1), // xn = RMSNorm(dx) (input_layernorm) + b.wq.record(r, d.xn.b, d.q.b, 1), // q/k/v = xn·Wq/Wk/Wv — NoPE, so no RoPE between here and the cache + b.wk.record(r, d.xn.b, d.k.b, 1), + b.wv.record(r, d.xn.b, d.v_.b, 1), + r.Blit(d.k.b, 0, b.kC, pos*kvDim, kvDim), // append this token's K/V rows to the cache + r.Blit(d.v_.b, 0, b.vC, pos*kvDim, kvDim), + d.recordMHA(r, d.q.b, b.kC, b.vC, d.attn.b, 1, pos+1), + d.recordOProj(r, b, 1), // dx += attn·Wo + ) +} + +// allocMambaScratch allocates the rows==1 Mamba decode scratch. Mamba has no attention or KV cache, +// so it skips the q/k/v/attn/qkv/gate/up buffers allocScratch would make; it keeps only the residual +// (dx), the pre-norm output (xn), the down-projection scratch (mo), the logits, and the per-block +// intermediates. Everything is one row wide — the sequential recurrence never batches. +func (d *Decoder) allocMambaScratch(mk func(data []float32) *bufSlot) { + dI, N, dtR := d.dInner, d.mambaN, d.dtRank + d.dx = mk(make([]float32, d.d)) + d.xn = mk(make([]float32, d.d)) + d.mo = mk(make([]float32, d.d)) + d.logits = mk(make([]float32, d.v)) + d.mbXin = mk(make([]float32, dI)) + d.mbZ = mk(make([]float32, dI)) + d.mbXc = mk(make([]float32, dI)) + d.mbDtLow = mk(make([]float32, dtR)) + d.mbDelta = mk(make([]float32, dI)) + d.mbB = mk(make([]float32, N)) + d.mbC = mk(make([]float32, N)) + d.mbY = mk(make([]float32, dI)) +} + +// recordMamba2Block records ONE Mamba-2 SSD block for a single decode token (rows==1), leaving +// dx += mixer(RMSNorm(dx)). It mirrors nlp.Mamba2Mixer.forward: fused in_proj → split [z | xBC | dt] +// → causal-conv step + SiLU over conv_dim → split [x | B | C] → Δ=softplus(dt+dt_bias) → the SSD scan +// (cu_ssd_step: scalar per-head decay a=exp(Δ·A), grouped B/C, +D·x skip) → gated RMSNorm norm(y·SiLU(z)) +// → out_proj onto the residual. A=−exp(A_log) precomputed on the host. Like Mamba-1 it is rows==1 +// (sequential recurrence). Splits are contiguous rows==1 sub-ranges, extracted with Blit. +func (d *Decoder) recordMamba2Block(r recorder, b block) error { + I, CD, H2 := d.m2Inter, d.m2CD, d.m2H + gN := d.m2G * d.m2N + // pre-norm → xn = RMSNorm(dx); fused in_proj → proj = [z(I) | xBC(CD) | dt(H2)]. + e := d.norm(r, d.dx.b, b.gAttn, nil, d.xn.b, 1) // RMSNorm (lnBias false) + e = firstErr(e, + b.m2InProj.record(r, d.xn.b, d.m2Proj.b, 1), + r.Blit(d.m2Proj.b, 0, d.m2Z.b, 0, I), // gate z + r.Blit(d.m2Proj.b, I, d.m2XBC.b, 0, CD), // conv input xBC + r.Blit(d.m2Proj.b, I+CD, d.m2Dt.b, 0, H2), // Δ pre-activation + // local mixing over conv_dim: causal depthwise-conv step (state carried) + SiLU. + r.Conv1DStep(d.m2XBC.b, b.m2ConvW, b.m2ConvB, b.m2ConvState, d.m2Xc.b, CD, d.m2Conv), + r.Unary(d.m2Xc.b, d.m2Xc.b, unarySiLU), + // split conv output → x(value, I) | B(gN) | C(gN). + r.Blit(d.m2Xc.b, 0, d.m2X.b, 0, I), + r.Blit(d.m2Xc.b, I, d.m2B.b, 0, gN), + r.Blit(d.m2Xc.b, I+gN, d.m2C.b, 0, gN), + // Δ = softplus(dt + dt_bias) [num_heads]. + r.AddBias(d.m2Dt.b, b.m2DtBias, d.m2Dt.b, 1, H2), + r.Unary(d.m2Dt.b, d.m2Dt.b, unarySoftplus), + // SSD scan: scalar per-head decay, grouped B/C, +D·x skip → y[I]. + r.SSDStep(d.m2X.b, d.m2Dt.b, b.m2A, d.m2B.b, d.m2C.b, b.m2Dskip, b.m2SsdState, d.m2Y.b, H2, d.m2P, d.m2G, d.m2N), + // gated RMSNorm: norm(y · SiLU(z)) · normW, over the intermediate width. + r.Unary(d.m2Z.b, d.m2Z.b, unarySiLU), + r.Binary(d.m2Y.b, d.m2Z.b, d.m2Y.b, binaryMul), + r.RMSNorm(d.m2Y.b, b.m2NormW, d.m2Y.b, 1, I, d.eps), + // out_proj onto the residual: dx += y · OutProj. + b.m2OutProj.recordAdd(r, d.m2Y.b, d.mo.b, d.dx.b, 1, d.d), + ) + return e +} + +// resetMamba2State zeroes every block's conv-window and SSD state (fresh sequence, pos==0). +func (d *Decoder) resetMamba2State() error { + zc := make([]float32, d.m2CD*(d.m2Conv-1)) + zs := make([]float32, d.m2H*d.m2N*d.m2P) + for _, b := range d.blocks { + if e := b.m2ConvState.UploadF32(zc); e != nil { + return e + } + if e := b.m2SsdState.UploadF32(zs); e != nil { + return e + } + } + return nil +} + +// allocMamba2Scratch allocates the rows==1 Mamba-2 decode scratch (no attention/KV buffers). +func (d *Decoder) allocMamba2Scratch(mk func(data []float32) *bufSlot) { + I, CD, H2 := d.m2Inter, d.m2CD, d.m2H + gN := d.m2G * d.m2N + projSize := I + CD + H2 + d.dx = mk(make([]float32, d.d)) + d.xn = mk(make([]float32, d.d)) + d.mo = mk(make([]float32, d.d)) + d.logits = mk(make([]float32, d.v)) + d.m2Proj = mk(make([]float32, projSize)) + d.m2Z = mk(make([]float32, I)) + d.m2XBC = mk(make([]float32, CD)) + d.m2Dt = mk(make([]float32, H2)) + d.m2Xc = mk(make([]float32, CD)) + d.m2X = mk(make([]float32, I)) + d.m2B = mk(make([]float32, gN)) + d.m2C = mk(make([]float32, gN)) + d.m2Y = mk(make([]float32, I)) +} + +// rwMix records the RWKV token-shift interpolation μ⊙x + (1−μ)⊙shift into out (rows==1). The +// reference forms it as shift + μ⊙(x−shift); with only add/mul recorder ops we use the algebraically +// identical μ⊙x + (1−μ)⊙shift, so mu and its precomputed complement omu are both uploaded. +func (d *Decoder) rwMixInto(r recorder, x, shift, mu, omu, out buffer) error { + return firstErr( + r.Binary(x, mu, d.rwT1.b, binaryMul), // μ⊙x + r.Binary(shift, omu, d.rwT2.b, binaryMul), // (1−μ)⊙shift + r.Binary(d.rwT1.b, d.rwT2.b, out, binaryAdd), + ) +} + +// recordRWKVBlock records ONE RWKV-4 block for a single decode token (rows==1), updating dx in place +// through both residuals. It mirrors nn.RWKVBlock.Step: a WKV TIME-MIX sublayer (LN1 → token-shift → +// r/k/v projections → the stabilized WKV recurrence via cu_wkv_step → receptance gate σ(r)⊙wkv → +// output projection → residual) and a gated squared-ReLU CHANNEL-MIX (LN2 → token-shift → σ(CWr·) +// gate over relu(CWk·)²·CWv → residual). No attention, no KV cache — the per-block token-shift and +// WKV states (rwPrevTM/rwPrevCM/rwAA/rwBB/rwPP) carry the whole history in O(1). Time-mix uses +// wq/wk/wv/wo as Wr/Wk/Wv/Wo; channel-mix uses wG/wD as CWk/CWv. +func (d *Decoder) recordRWKVBlock(r recorder, b block) error { + // TIME-MIX: xn = LN1(dx); r=σ(Wr·mix_r), k=Wk·mix_k, v=Wv·mix_v. + e := r.LayerNorm(d.dx.b, b.gAttn, b.bAttn, d.xn.b, 1, d.d, d.eps) + e = firstErr(e, + d.rwMixInto(r, d.xn.b, b.rwPrevTM, b.rwMuR, b.rwOmR, d.rwMix.b), + b.wq.record(r, d.rwMix.b, d.rwR.b, 1), + r.Unary(d.rwR.b, d.rwR.b, unarySigmoid), + d.rwMixInto(r, d.xn.b, b.rwPrevTM, b.rwMuK, b.rwOmK, d.rwMix.b), + b.wk.record(r, d.rwMix.b, d.rwK.b, 1), + d.rwMixInto(r, d.xn.b, b.rwPrevTM, b.rwMuV, b.rwOmV, d.rwMix.b), + b.wv.record(r, d.rwMix.b, d.rwV.b, 1), + // WKV recurrence (state carried), then receptance gate and output projection onto the residual. + r.WKVStep(d.rwK.b, d.rwV.b, b.rwW, b.rwU, b.rwAA, b.rwBB, b.rwPP, d.rwWkv.b, d.d), + r.Binary(d.rwWkv.b, d.rwR.b, d.rwWkv.b, binaryMul), // σ(r)⊙wkv (o==a: Binary clobbers b, not a) + b.wo.record(r, d.rwWkv.b, d.rwO.b, 1), + r.Binary(d.dx.b, d.rwO.b, d.dx.b, binaryAdd), // residual + r.Blit(d.xn.b, 0, b.rwPrevTM, 0, d.d), // prevTM ← xn (token-shift state) + ) + // CHANNEL-MIX: yn = LN2(dx); cr=σ(CWr·mix_r); ck=relu(CWk·mix_k)²; cv=CWv·ck; dx += cr⊙cv. + e = firstErr(e, + r.LayerNorm(d.dx.b, b.gFFN, b.bFFN, d.xn.b, 1, d.d, d.eps), + d.rwMixInto(r, d.xn.b, b.rwPrevCM, b.rwCMuR, b.rwOmCR, d.rwMix.b), + b.rwCWr.record(r, d.rwMix.b, d.rwCr.b, 1), + r.Unary(d.rwCr.b, d.rwCr.b, unarySigmoid), + d.rwMixInto(r, d.xn.b, b.rwPrevCM, b.rwCMuK, b.rwOmCK, d.rwMix.b), + b.wG.record(r, d.rwMix.b, d.rwCk.b, 1), // CWk: Dim→Hidden + r.Unary(d.rwCk.b, d.rwCk.b, unaryReLU2), // relu(·)² + b.wD.record(r, d.rwCk.b, d.rwCv.b, 1), // CWv: Hidden→Dim + r.Binary(d.rwCv.b, d.rwCr.b, d.rwCv.b, binaryMul), // σ(CWr)⊙cv (o==a: Binary clobbers b, not a) + r.Binary(d.dx.b, d.rwCv.b, d.dx.b, binaryAdd), // residual + r.Blit(d.xn.b, 0, b.rwPrevCM, 0, d.d), // prevCM ← yn + ) + return e +} + +// resetRWKVState zeroes every block's token-shift + WKV state for a fresh sequence (pos==0): prevTM/ +// prevCM/aa/bb start at 0 and pp at −1e38 (the pre-first-token WKV running-max, matching OpWKV). +func (d *Decoder) resetRWKVState() error { + zero := make([]float32, d.d) + ninf := make([]float32, d.d) + for i := range ninf { + ninf[i] = -1e38 + } + for _, b := range d.blocks { + for _, st := range []buffer{b.rwPrevTM, b.rwPrevCM, b.rwAA, b.rwBB} { + if e := st.UploadF32(zero); e != nil { + return e + } + } + if e := b.rwPP.UploadF32(ninf); e != nil { + return e + } + } + return nil +} + +// allocRWKVScratch allocates the rows==1 RWKV decode scratch (no attention/KV buffers). +func (d *Decoder) allocRWKVScratch(mk func(data []float32) *bufSlot) { + d.dx = mk(make([]float32, d.d)) + d.xn = mk(make([]float32, d.d)) + d.logits = mk(make([]float32, d.v)) + d.rwMix = mk(make([]float32, d.d)) + d.rwR = mk(make([]float32, d.d)) + d.rwK = mk(make([]float32, d.d)) + d.rwV = mk(make([]float32, d.d)) + d.rwWkv = mk(make([]float32, d.d)) + d.rwO = mk(make([]float32, d.d)) + d.rwT1 = mk(make([]float32, d.d)) + d.rwT2 = mk(make([]float32, d.d)) + d.rwCr = mk(make([]float32, d.d)) + d.rwCk = mk(make([]float32, d.rwHidden)) + d.rwCv = mk(make([]float32, d.d)) +} + +func (d *Decoder) recordFFN(r recorder, b block, in buffer, rows int) error { + if d.ffnGELU || d.ffnReLU2 { + act := unaryGELU // 2-layer MLP activation: GELU (GPT-NeoX/Phi/StarCoder2) or squared ReLU (Nemotron) + if d.ffnReLU2 { + act = unaryReLU2 + } + e := b.wG.record(r, in, d.gate.b, rows) // fc = in·Wfc + if b.fcBias != nil { + e = firstErr(e, r.AddBias(d.gate.b, b.fcBias, d.gate.b, rows, d.hidden)) + } + e = firstErr(e, + r.Unary(d.gate.b, d.gate.b, act), // act(fc) + d.recordDownProj(r, b, rows), // dx += act(fc)·Wproj + ) + if b.projBias != nil { + e = firstErr(e, r.AddBias(d.dx.b, b.projBias, d.dx.b, rows, d.d)) + } + return e + } + if d.ffnGEGLU { + // GeGLU (Gemma): down(GELU(gate)⊙up). Same 3-matrix gated shape as SwiGLU, but the gate + // activation is GELU and the elementwise product is a separate mul (no fused silu·b op). + if b.wGU != nil { + // Column-fused: one GEMV into gu = [·, 2·hidden], then GeGLU over its halves (GELU⊙up) into + // gate — one launch instead of two GEMVs + Unary GELU + Binary mul. Bit-identical. + e := b.wGU.record(r, in, d.gu.b, rows) + if gh, ok := r.(geGLUHalvesRecorder); ok { + e = firstErr(e, gh.GeGLUHalves(d.gu.b, d.gate.b, rows, d.hidden)) + } else { + e = firstErr(e, + r.Copy2D(d.gu.b, 0, 2*d.hidden, d.gate.b, 0, d.hidden, rows, d.hidden), + r.Copy2D(d.gu.b, d.hidden, 2*d.hidden, d.up.b, 0, d.hidden, rows, d.hidden), + r.Unary(d.gate.b, d.gate.b, unaryGELU), + d.binElem(r, d.gate.b, d.up.b, d.gate.b, binaryMul, rows, d.hidden), + ) + } + return firstErr(e, d.recordDownProj(r, b, rows)) + } + return firstErr( + b.wG.record(r, in, d.gate.b, rows), // gate = in·Wgate + b.wU.record(r, in, d.up.b, rows), // up = in·Wup + r.Unary(d.gate.b, d.gate.b, unaryGELU), // GELU(gate) + d.binElem(r, d.gate.b, d.up.b, d.gate.b, binaryMul, rows, d.hidden), // GELU(gate)⊙up + d.recordDownProj(r, b, rows), // dx += (·)·Wdown + ) + } + if b.wGU != nil { + // Column-fused ffn_gate|ffn_up: one GEMV into gu = [·, 2·hidden], then SwiGLU over its halves + // into gate — one launch instead of two GEMVs, bit-identical (each output column quantizes + // independently, so the fused weight equals the two separate ones concatenated). + e := b.wGU.record(r, in, d.gu.b, rows) + if sh, ok := r.(swiGLUHalvesRecorder); ok { + e = firstErr(e, sh.SwiGLUHalves(d.gu.b, d.gate.b, rows, d.hidden)) + } else { + // Backend without the fused op: split the gu halves out (strided) and SwiGLU as usual. + e = firstErr(e, + r.Copy2D(d.gu.b, 0, 2*d.hidden, d.gate.b, 0, d.hidden, rows, d.hidden), + r.Copy2D(d.gu.b, d.hidden, 2*d.hidden, d.up.b, 0, d.hidden, rows, d.hidden), + d.binElem(r, d.gate.b, d.up.b, d.gate.b, binarySwiGLU, rows, d.hidden), + ) + } + return firstErr(e, d.recordDownProj(r, b, rows)) // dx += swiglu·Wdown + } + e := b.wG.record(r, in, d.gate.b, rows) + e = firstErr(e, b.wU.record(r, in, d.up.b, rows)) + e = firstErr(e, recordBarrier(r)) + e = firstErr(e, d.binElem(r, d.gate.b, d.up.b, d.gate.b, binarySwiGLU, rows, d.hidden)) + e = firstErr(e, recordBarrier(r)) + return firstErr(e, d.recordDownProj(r, b, rows)) // dx += swiglu·Wdown +} + +// binElem records an elementwise op bounded to the rows*width elements actually in play. Decoder +// buffers are allocated for the whole context, so the unbounded form touches ctx*width elements for +// a single decoded row. Recorders without the capability fall back to the unbounded call, which is +// correct — the tail of the buffer holds stale values that no consumer reads — just far slower. +func (d *Decoder) binElem(r recorder, a, b, o buffer, op, rows, width int) error { + if bn, ok := r.(binaryNRecorder); ok && rows > 0 && width > 0 { + return bn.BinaryN(a, b, o, op, rows*width) + } + return r.Binary(a, b, o, op) +} + +// norm records the pre-sublayer normalization: LayerNorm-with-bias when lnBias (StableLM/Phi), +// else RMSNorm (Llama). gamma is the weight; beta is the LayerNorm bias (ignored under RMSNorm). +func (d *Decoder) norm(r recorder, x, gamma, beta, o buffer, rows int) error { + if d.lnBias { + return r.LayerNorm(x, gamma, beta, o, rows, d.d, d.eps) + } + return r.RMSNorm(x, gamma, o, rows, d.d, d.eps) +} + +// bFinalBeta is the final-norm LayerNorm bias, or nil under RMSNorm (bFinal unset). +func (d *Decoder) bFinalBeta() buffer { + if d.bFinal != nil { + return d.bFinal.b + } + return nil +} + +func (d *Decoder) ropeQK(r recorder, qkv, inv buffer, seq, stride, pos int) error { + if d.noRope { // MPT: position enters through the ALiBi attention bias, not rotary + return nil + } + if d.rotaryDim < d.dk { + return r.RoPEPartialPair(qkv, inv, seq, stride, d.h, 0, d.kvH, d.d, d.dk, d.rotaryDim, pos, d.posDiv) + } + return r.RoPEPair(qkv, inv, seq, stride, d.h, 0, d.kvH, d.d, d.dk, d.half, pos, d.posDiv) +} + +// allocScratch uploads the RoPE frequencies and allocates the per-step scratch buffers. They are +// sized for up to maxLen rows so StepN can process a whole prompt (or a speculative draft window) +// in one recorded step (§T418) — a few MB at typical configs. +func (d *Decoder) allocScratch(mk func(data []float32) *bufSlot) { + c := d.maxLen + d.dinv = mk(d.invHost) + d.dx = mk(make([]float32, c*d.d)) + d.xn = mk(make([]float32, c*d.d)) + d.xn2 = mk(make([]float32, c*d.d)) + d.q = mk(make([]float32, c*d.qDim)) + d.k = mk(make([]float32, c*d.kvDim)) + d.v_ = mk(make([]float32, c*d.kvDim)) + d.qkv = mk(make([]float32, c*(d.qDim+2*d.kvDim))) + d.attn = mk(make([]float32, c*d.qDim)) + d.ao = mk(make([]float32, c*d.d)) + d.gate = mk(make([]float32, c*d.hidden)) + d.up = mk(make([]float32, c*d.hidden)) + if d.ops.fusedGateUp { + d.gu = mk(make([]float32, c*2*d.hidden)) // fused gate|up GEMV output for the SwiGLUHalves path + } + d.mo = mk(make([]float32, c*d.d)) + d.logitsRows = 1 + if d.ops.eagerFullLogits { + d.logitsRows = c + } + d.logits = mk(make([]float32, d.logitsRows*d.v)) + if d.moe { + d.moeGate = mk(make([]float32, c*d.nExperts)) + d.moeW = mk(make([]float32, c*d.nExperts)) + d.moeCol = mk(make([]float32, c)) + } +} + +// mkBuf returns a bufSlot allocator that records the first error and tracks buffers for Release. +func (d *Decoder) mkBuf(err *error) func(data []float32) *bufSlot { + return func(data []float32) *bufSlot { + if *err != nil { + return &bufSlot{} + } + b, e := d.ops.newBuffer(data) + if e != nil { + *err = e + return &bufSlot{} + } + d.all = append(d.all, b) + return &bufSlot{b: b} + } +} + +// mkLin is the shared quantize-aware projection builder for the per-arch constructors: resident Q8_0 +// when the backend set ops.quantizeF32 (a NewXQ8CUDA entry point), else the f32 weight — byte-identical +// to the old inline `lin` closures when the hook is nil, so every existing f32 constructor is unchanged. +// One helper lets an arch gain Q8 decode by adding a hook-setting constructor, no per-arch code change. +func (d *Decoder) mkLin(mk func([]float32) *bufSlot, ops backendOps, errp *error) func(*tensor.Tensor) linear { + return func(w *tensor.Tensor) linear { + in, out := w.Shape()[0], w.Shape()[1] + if ops.quantizeF32 == nil { + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } + qw, qe := ops.quantizeF32(w) + if qe != nil { + if *errp == nil { + *errp = qe + } + return f32Linear{k: in, n: out} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } +} + +// f32Lin is the always-f32 projection builder, ignoring any Q8 hook. Used for weights whose +// quantization would be incorrect even under a NewXQ8CUDA path — notably the MoE router: Q8 rounding in +// the [dim,E] router logits could flip a top-k expert selection and change the output discontinuously, +// so routing stays exact while the (bandwidth-dominant) expert matrices go Q8. +func (d *Decoder) f32Lin(mk func([]float32) *bufSlot) func(*tensor.Tensor) linear { + return func(w *tensor.Tensor) linear { + in, out := w.Shape()[0], w.Shape()[1] + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } +} + +// mkLinS is mkLin for a weight pre-scaled by s (a per-element scalar folded into the upload — Granite's +// ResidualMult / logit-scale, Cohere's logit_scale). Q8_0 when ops.quantizeF32 is set (the scaled weight +// is quantized POST-scale), else f32 — byte-identical to the old inline f32-only linS when the hook is nil. +// Fixes a latent partial-Q8 bug: the per-arch f32-only linS left scaled lm_head / Wo / Wdown weights f32 +// even under a NewXQ8CUDA hook. +func (d *Decoder) mkLinS(mk func([]float32) *bufSlot, ops backendOps, errp *error) func(*tensor.Tensor, float32) linear { + return func(w *tensor.Tensor, s float32) linear { + in, out := w.Shape()[0], w.Shape()[1] + f := flat2D(w) + if s != 1 { + for i := range f { + f[i] *= s + } + } + if ops.quantizeF32 == nil { + return f32Linear{w: mk(f).b, k: in, n: out} + } + t := tensor.New(tensor.F32, tensor.Shape{in, out}) + copy(t.Storage().F32(), f) + qw, qe := ops.quantizeF32(t) + if qe != nil { + if *errp == nil { + *errp = qe + } + return f32Linear{k: in, n: out} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } +} + +// mkFused is mkLin for the fused QKV weight (§T613): the three [in,out] projections concatenated per +// input row into one [in, nq+nk+nv] matrix, then Q8_0-quantized whole (or f32). cu_qmatmul_q8 is a GEMM +// (rows>1), so batched-prefill StepN works; qkv bias / QK-norm / RoPE stay f32 and apply after. +func (d *Decoder) mkFused(mk func([]float32) *bufSlot, ops backendOps, errp *error) func(wq, wk, wv *tensor.Tensor) linear { + return func(wq, wk, wv *tensor.Tensor) linear { + in := wq.Shape()[0] + nq, nk, nv := wq.Shape()[1], wk.Shape()[1], wv.Shape()[1] + fq, fk, fv := flat2D(wq), flat2D(wk), flat2D(wv) + nt := nq + nk + nv + w := make([]float32, in*nt) + for i := range in { + row := w[i*nt : (i+1)*nt] + copy(row[:nq], fq[i*nq:(i+1)*nq]) + copy(row[nq:nq+nk], fk[i*nk:(i+1)*nk]) + copy(row[nq+nk:], fv[i*nv:(i+1)*nv]) + } + if ops.quantizeF32 == nil { + return f32Linear{w: mk(w).b, k: in, n: nt} + } + t := tensor.New(tensor.F32, tensor.Shape{in, nt}) + copy(t.Storage().F32(), w) + qw, qe := ops.quantizeF32(t) + if qe != nil { + if *errp == nil { + *errp = qe + } + return f32Linear{k: in, n: nt} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } +} + +// mkFused2 is mkFused for the fused gate|up weight: two [in,out] projections concatenated per input +// row into one [in, ng+nu] matrix, then quantized whole (or f32) — identical per-column quantization +// to the separate weights (each output column quantizes over its own K independently), so the fused +// GEMV output is bit-identical to two separate GEMVs. Used only when ops.fusedGateUp (CUDA SwiGLU). +func (d *Decoder) mkFused2(mk func([]float32) *bufSlot, ops backendOps, errp *error) func(wg, wu *tensor.Tensor) linear { + return func(wg, wu *tensor.Tensor) linear { + in := wg.Shape()[0] + ng, nu := wg.Shape()[1], wu.Shape()[1] + fg, fu := flat2D(wg), flat2D(wu) + nt := ng + nu + w := make([]float32, in*nt) + for i := range in { + row := w[i*nt : (i+1)*nt] + copy(row[:ng], fg[i*ng:(i+1)*ng]) + copy(row[ng:], fu[i*nu:(i+1)*nu]) + } + if ops.quantizeF32 == nil { + return f32Linear{w: mk(w).b, k: in, n: nt} + } + t := tensor.New(tensor.F32, tensor.Shape{in, nt}) + copy(t.Storage().F32(), w) + qw, qe := ops.quantizeF32(t) + if qe != nil { + if *errp == nil { + *errp = qe + } + return f32Linear{k: in, n: nt} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } +} + +// gateUpBuilder returns a helper that, per block, either column-fuses ffn_gate|ffn_up into one wGU +// weight (when ops.fusedGateUp — the CUDA SwiGLU/GeGLU path) or builds the separate wG/wU pair. Lets +// every SwiGLU/GeGLU builder opt into the fused FFN with one line, keeping non-CUDA backends identical. +func (d *Decoder) gateUpBuilder(mk func([]float32) *bufSlot, ops backendOps, errp *error) func(wg, wu *tensor.Tensor) (gu, g, u linear) { + lin := d.mkLin(mk, ops, errp) + fuseGU := d.mkFused2(mk, ops, errp) + return func(wg, wu *tensor.Tensor) (gu, g, u linear) { + if ops.fusedGateUp { + return fuseGU(wg, wu), nil, nil + } + return nil, lin(wg), lin(wu) + } +} + +// newDecoder uploads m's f32 weights via ops into device buffers and prepares a KV cache up to Ctx. +func newDecoder(m *nlp.Llama, ops backendOps) (*Decoder, error) { + cfg := m.Config + // A blockless model is an ERROR, never a legitimate degenerate case. The upload loop below is + // `for _, b := range m.Blocks`, so with no blocks it simply never runs and Step falls through + // embedding → final norm → lm_head: plausible logits, a full Generate sequence, err == nil + // everywhere, and the entire transformer stack silently skipped — the confident wrong answer + // §V29 forbids. Same guard as newDeepSeekV2Decoder / newMambaDecoder / newRWKVDecoder and the + // MoE builders; newDecoder (behind New, NewCUDA, NewVulkan, NewLlamaQ8CUDA, NewLlamaQ4KCUDA and + // the Qwen2/Qwen3/Phi3/Granite constructors) was the only sibling missing it. + if len(m.Blocks) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Llama has no blocks", ops.name) + } + // The same silence one step milder: a config claiming more (or fewer) layers than the checkpoint + // carries decodes the blocks it has and says nothing about the ones it does not. Nothing here + // reads cfg.Layers, so a config that leaves it unset (0) is not a claim and stays accepted — + // only an explicit disagreement is rejected. + if cfg.Layers > 0 && cfg.Layers != len(m.Blocks) { + return nil, fmt.Errorf("llamagpu(%s): config claims %d layers but model has %d blocks", ops.name, cfg.Layers, len(m.Blocks)) + } + d, derr := newDecoderCommon(cfg, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + + var err error + mk := d.mkBuf(&err) + // qOrF32 is the projection core: resident Q8_0 when the backend set ops.quantizeF32 (the + // NewLlamaQ8CUDA / NewQwen2Q8CUDA / NewQwen3Q8CUDA / NewPhi3Q8CUDA family), else the f32 weight. + // EVERY projection — fused QKV, o_proj, gate/up/down, lm_head — flows through it, so an f32 + // checkpoint gets a direct Q8 decode path (~4× less projection weight bandwidth) without a ggml + // conversion — and it reaches the Qwen2 (qkv-bias) / Qwen3 (QK-norm) variants that nlp.QuantLlama + // (Llama-only) cannot represent. Biases and QK-norm gains stay f32 and apply after the Q8 matmul. + qOrF32 := func(f []float32, in, out int) linear { + if ops.quantizeF32 == nil { + return f32Linear{w: mk(f).b, k: in, n: out} + } + t := tensor.New(tensor.F32, tensor.Shape{in, out}) + copy(t.Storage().F32(), f) + qw, qe := ops.quantizeF32(t) + if qe != nil { + if err == nil { + err = qe + } + return f32Linear{k: in, n: out} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } + lin := func(w *tensor.Tensor) linear { // Q8_0 (quantized build) or f32 [in,out] device weight + in, out := w.Shape()[0], w.Shape()[1] + return qOrF32(flat2D(w), in, out) + } + // IBM Granite config scalars (all identity when 0 or 1, so no-ops for Llama/Qwen/Mistral). Each + // folds into the upload rather than a runtime op: AttentionMult overrides the pre-softmax scale, + // EmbeddingMult scales gathered embeddings (d.embMult), ResidualMult scales each sublayer output + // (baked into Wo and Wdown, since their matmul IS the residual-add epilogue — Granite carries no + // projection bias, so nothing else in the add needs scaling), and LogitsScale divides the logits + // (baked as 1/LogitsScale into the untied lm_head). Separate uploads, so tied embed/lm_head is fine. + if cfg.AttentionMult != 0 { + d.scale = float32(cfg.AttentionMult) + } + if cfg.EmbeddingMult != 0 && cfg.EmbeddingMult != 1 { + d.embMult = float32(cfg.EmbeddingMult) + } + resMult := float32(1) + if cfg.ResidualMult != 0 && cfg.ResidualMult != 1 { + resMult = float32(cfg.ResidualMult) + } + outScale := float32(1) + if cfg.LogitsScale != 0 && cfg.LogitsScale != 1 { + outScale = float32(1 / cfg.LogitsScale) + } + linS := func(w *tensor.Tensor, s float32) linear { // weight with every element pre-scaled by s (then Q8 ∨ f32) + in, out := w.Shape()[0], w.Shape()[1] + f := flat2D(w) + if s != 1 { + for i := range f { + f[i] *= s + } + } + return qOrF32(f, in, out) + } + // fused QKV weight (§T613): weights are [in,out] with out along the row, so the fusion + // concatenates the three output bands PER INPUT ROW — one [in, D+2·kvDim] matrix whose + // product row is [q | k | v]. The separate wq/wk/wv uploads are dropped (no dup storage). + fused := func(wq, wk, wv *tensor.Tensor) linear { + in := wq.Shape()[0] + nq, nk, nv := wq.Shape()[1], wk.Shape()[1], wv.Shape()[1] + fq, fk, fv := flat2D(wq), flat2D(wk), flat2D(wv) + nt := nq + nk + nv + w := make([]float32, in*nt) + for i := range in { + row := w[i*nt : (i+1)*nt] + copy(row[:nq], fq[i*nq:(i+1)*nq]) + copy(row[nq:nq+nk], fk[i*nk:(i+1)*nk]) + copy(row[nq+nk:], fv[i*nv:(i+1)*nv]) + } + return qOrF32(w, in, nt) + } + // fused q/k/v bias [Bq | Bk | Bv] = [D+2·kvDim], for the Qwen2/Qwen2.5 family (q/k/v carry a + // projection bias, o_proj does not); nil for Llama/Mistral, which leaves recordQKVProj's + // AddBias unrecorded so those models stay byte-identical. + fusedBias := func(bq, bk, bv *tensor.Tensor) buffer { + if bq == nil && bk == nil && bv == nil { + return nil + } + fb := append(append(append([]float32{}, flat1D(bq)...), flat1D(bk)...), flat1D(bv)...) + return mk(fb).b + } + // per-head query/key RMSNorm gain [dk] (Qwen3); nil for Llama/Qwen2. Presence flips d.qkNorm, + // which arms recordQKNorm between the QKV projection and RoPE. + qkGain := func(n *nn.RMSNorm) buffer { + if n == nil { + return nil + } + d.qkNorm = true + return mk(flat1D(n.Gamma)).b + } + // gateUp column-fuses ffn_gate|ffn_up into one wGU weight when the backend can SwiGLU-over-halves + // (CUDA), else keeps the separate wG/wU pair. Same quantization path as the QKV fusion (qOrF32). + fuseGU := d.mkFused2(mk, ops, &err) + gateUp := func(wg, wu *tensor.Tensor) (gu, g, u linear) { + if ops.fusedGateUp { + return fuseGU(wg, wu), nil, nil + } + return nil, lin(wg), lin(wu) + } + //perfscan:ignore PS3032 model-build weight-upload loop over blocks, one-time + for _, b := range m.Blocks { + wgu, wg, wu := gateUp(b.FFN.Wgate, b.FFN.Wup) + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), qkvBias: fusedBias(b.Bq, b.Bk, b.Bv), wo: linS(b.Wo, resMult), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + qN: qkGain(b.QNorm), kN: qkGain(b.KNorm), + wGU: wgu, wG: wg, wU: wu, wD: linS(b.FFN.Wdown, resMult), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = linS(m.Out, outScale) // LogitsScale folded into the lm_head (1/LogitsScale) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newStableLMDecoder builds a device Decoder for an nlp.StableLM: the Llama-shaped SwiGLU/GQA +// core, but with LayerNorm-with-bias (lnBias) and PARTIAL rotary (rotaryDim < dk). StableLM has +// no attention bias and its rotary is split-half (matches OpRoPE), so weights upload directly — +// only the norms carry a bias and the rope is partial. Currently cuda-only (metal/vulkan +// RoPEPartialPair are unimplemented). §T757-class new-architecture GPU decode. +func newStableLMDecoder(m *nlp.StableLM, ops backendOps) (*Decoder, error) { + cfg := m.Config + // StableLMConfig's kvHeads()/headDim()/rotaryDim() are unexported, so replicate here. + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.HeadDim + if headDim <= 0 { + headDim = cfg.Dim / cfg.Heads + } + rotaryDim := cfg.RotaryDim + if rotaryDim <= 0 { + pct := cfg.RotaryPct + if pct == 0 { + pct = 0.25 + } + rotaryDim = int(float64(headDim) * pct) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + // LayerNorm-bias + partial rotary: rebuild the inverse-frequency table over rotaryDim + // (newDecoderCommon sized it to dk for full rope). + d.lnBias = true + d.rotaryDim = rotaryDim + if d.rotaryDim <= 0 || d.rotaryDim > d.dk || d.rotaryDim%2 != 0 { + return nil, fmt.Errorf("llamagpu(%s): StableLM rotaryDim %d invalid for headDim %d", ops.name, d.rotaryDim, d.dk) + } + invF64, posDiv64 := backend.RoPEFreqs(d.rotaryDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, d.rotaryDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + //perfscan:ignore PS3032 newStableLMDecoder block-upload, one-time model build + for _, b := range m.Blocks { + wgu, wg, wu := gateUp(b.FFN.Wgate, b.FFN.Wup) + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + gFFN: mk(flat1D(b.PostAttnNorm.Gamma)).b, bFFN: mk(flat1D(b.PostAttnNorm.Beta)).b, + wGU: wgu, wG: wg, wU: wu, wD: lin(b.FFN.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.bFinal = mk(flat1D(m.Norm.Beta)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newCohereDecoder builds a device Decoder for an nlp.Cohere (Command-R). Cohere composes existing +// generalizations: ONE-norm parallel residual (parallelRes — input_layernorm feeds both attention +// and FFN, summed onto x0, exactly like Phi), weight-only mean-centered LayerNorm (lnBias with a +// zero β), SwiGLU, GQA and FULL rope. Its interleaved (GPT-J) rotary is already baked into a q/k +// weight-row permutation by CohereFromHF, so standard split-half RoPE reproduces it. The only extra +// is logit_scale: the tied lm_head's logits are multiplied by a constant, folded into the Out +// upload. cuda-only (parallel-residual + LayerNorm reuse the Phi plumbing). +func newCohereDecoder(m *nlp.Cohere, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + if cfg.HeadDim != 0 && cfg.HeadDim != cfg.Dim/cfg.Heads { + return nil, fmt.Errorf("llamagpu(%s): Cohere decoupled head_dim %d (≠ dim/heads %d) not supported on the GPU decoder yet", ops.name, cfg.HeadDim, cfg.Dim/cfg.Heads) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // weight-only mean-centered LayerNorm (β is a zero vector) + d.parallelRes = true // one-norm parallel residual: input_layernorm feeds attn AND FFN + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + logitScale := float32(1) + if cfg.LogitScale != 0 && cfg.LogitScale != 1 { + logitScale = float32(cfg.LogitScale) + } + linS := d.mkLinS(mk, ops, &err) + //perfscan:ignore PS3032 newCohereDecoder block-upload, one-time model build + for _, b := range m.Blocks { + wgu, wg, wu := gateUp(b.FFN.Wgate, b.FFN.Wup) + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + // gFFN/bFFN unused: parallel one-norm reuses the attn norm (d.xn) for the FFN. + wGU: wgu, wG: wg, wU: wu, wD: lin(b.FFN.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.bFinal = mk(flat1D(m.Norm.Beta)) + d.out = linS(m.Out, logitScale) // logit_scale folded into the tied lm_head + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newNemotronDecoder builds a device Decoder for an nlp.Nemotron. Nemotron is a SEQUENTIAL +// two-norm decoder (input_layernorm feeds attention, post_attention_layernorm feeds the MLP, each +// on its own residual add) with three departures from Llama, all reusing existing plumbing: +// LayerNorm1P (an ordinary mean-centered LayerNorm — lnBias — whose γ=w+1 and β are folded in by +// NemotronFromHF), a squared-ReLU 2-layer MLP (ffnReLU2: down(relu²(up(x))), no gate, no bias), and +// PARTIAL rotary. cuda-only (partial rotary + the relu² unary are cuda-only). +func newNemotronDecoder(m *nlp.Nemotron, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.HeadDim + if headDim <= 0 { + headDim = cfg.Dim / cfg.Heads + } + if headDim != cfg.Dim/cfg.Heads { + return nil, fmt.Errorf("llamagpu(%s): Nemotron decoupled head_dim %d (≠ dim/heads %d) not supported on the GPU decoder yet", ops.name, headDim, cfg.Dim/cfg.Heads) + } + rotaryDim := cfg.RotaryDim + if rotaryDim <= 0 { + pct := cfg.RotaryPct + if pct == 0 { + pct = 0.5 + } + rotaryDim = int(float64(headDim) * pct) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // LayerNorm1P = mean-centered LayerNorm with a bias + d.ffnReLU2 = true // squared-ReLU 2-layer MLP + d.rotaryDim = rotaryDim + if d.rotaryDim <= 0 || d.rotaryDim > d.dk || d.rotaryDim%2 != 0 { + return nil, fmt.Errorf("llamagpu(%s): Nemotron rotaryDim %d invalid for headDim %d", ops.name, d.rotaryDim, d.dk) + } + invF64, posDiv64 := backend.RoPEFreqs(d.rotaryDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, d.rotaryDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + //perfscan:ignore PS3032 newNemotronDecoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + gFFN: mk(flat1D(b.PostAttnNorm.Gamma)).b, bFFN: mk(flat1D(b.PostAttnNorm.Beta)).b, + wG: lin(b.Wup), wD: lin(b.Wdown), // 2-layer relu² MLP: wG = up (fc), wD = down (proj); no gate + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.bFinal = mk(flat1D(m.Norm.Beta)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newGemmaDecoder builds a device Decoder for an nlp.Gemma (Gemma v1). Gemma is close to Llama — +// pre-norm RMSNorm, RoPE, GQA — with three departures, all reusing existing plumbing: the token +// embeddings are scaled by √dim right after lookup (d.embMult), RMSNorm uses (1+w) as the gain +// (folded in at load, so nn.RMSNorm applies directly), and the FFN is GeGLU rather than SwiGLU +// (ffnGEGLU — GELU gate + a plain elementwise product). The lm_head is tied to the embedding +// (Out = TokEmbᵀ). cuda-only for consistency with the rest of the new-arch family. +func newGemmaDecoder(m *nlp.Gemma, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.HeadDim + if headDim <= 0 { + headDim = cfg.Dim / cfg.Heads + } + // Gemma infers head_dim from q_proj (heads·head_dim rows) and it is DECOUPLED (head_dim need not + // equal dim/heads), so read it back from the loaded weights rather than trusting the config — + // GemmaFromHF may have left cfg.HeadDim at 0. + if qOut := m.Blocks[0].Wq.Shape()[1]; qOut%cfg.Heads == 0 { + headDim = qOut / cfg.Heads + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.FFN, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.ffnGEGLU = true // GELU-gated FFN + d.embMult = float32(math.Sqrt(float64(cfg.Dim))) // Gemma's √dim embedding normalizer + // Decoupled head_dim: override the geometry newDecoderCommon derived as dim/heads. qDim (the + // attention Q/O width) becomes h·head_dim, which for Gemma differs from the model dim d. + if headDim != d.dk { + d.dk = headDim + d.qDim = d.h * headDim + d.kvDim = d.kvH * headDim + d.half = headDim / 2 + d.rotaryDim = headDim + d.scale = float32(1.0 / math.Sqrt(float64(headDim))) + invF64, posDiv64 := backend.RoPEFreqs(headDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, headDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + } + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + // Column-fuse ffn_gate|ffn_up into wGU when the backend can GeGLU-over-halves (CUDA), else keep + // the separate wG/wU pair. GeGLU (Gemma) uses the same wGU layout as SwiGLU; only the halves op differs. + fuseGU := d.mkFused2(mk, ops, &err) + gateUp := func(wg, wu *tensor.Tensor) (gu, g, u linear) { + if ops.fusedGateUp { + return fuseGU(wg, wu), nil, nil + } + return nil, lin(wg), lin(wu) + } + //perfscan:ignore PS3032 newGemmaDecoder block-upload, one-time model build + for _, b := range m.Blocks { + wgu, wg, wu := gateUp(b.FFN.Wgate, b.FFN.Wup) + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + wGU: wgu, wG: wg, wU: wu, wD: lin(b.FFN.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + // Tied lm_head: Out = TokEmbᵀ. TokEmb is [vocab, dim]; the projection wants [dim, vocab]. + vocab, dim := m.TokEmb.Shape()[0], m.TokEmb.Shape()[1] + // Route the tied lm_head through the same quantize-aware builder so the Q8 constructors quantize it too + // — Gemma's vocab is huge (256k), so an f32 lm_head would dominate decode bandwidth and dilute Q8. The + // f32 path is byte-identical (mkLin with no hook == the old f32Linear{mk(outW)}). + outT := tensor.New(tensor.F32, tensor.Shape{dim, vocab}) + flat2DTInto(outT.Storage().F32()[:dim*vocab], m.TokEmb) + d.out = lin(outT) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newDeepSeekV2Decoder builds a device Decoder for an nlp.DeepSeekV2 — a stack of Multi-head Latent +// Attention blocks, each with a dense SwiGLU or a DeepSeekMoE sparse FFN (first_k_dense_replace +// split). MLA is the +// only new attention machinery (recordMLAAttention + rectangular MHARect); the FFN is a per-block +// dense SwiGLU or the DeepSeekMoE sparse FFN (raw-softmax·routedScale top-k routed experts + an +// ungated shared expert), dispatched by b.moeRouter. cuda-only (MHARect is cuda-only). +func newDeepSeekV2Decoder(m *nlp.DeepSeekV2, ops backendOps) (*Decoder, error) { + cfg := m.Config + if len(m.Blocks) == 0 { + return nil, fmt.Errorf("llamagpu(%s): DeepSeek-V2 has no blocks", ops.name) + } + qkHead := cfg.QKNope + cfg.QKRope + kvHead := cfg.QKNope + cfg.VHead + // hidden must cover the widest SwiGLU across every block (dense, routed-expert, shared) since the + // gate/up scratch is shared; nExperts/topK/routedScale come from the first MoE block if any. + hidden, nExperts, topK, hasMoE := 0, 0, 0, false + wider := func(sw *nn.SwiGLU) { + if sw != nil { + if h := sw.Wgate.Shape()[1]; h > hidden { + hidden = h + } + } + } + for _, b := range m.Blocks { + wider(b.Dense) + if b.MoE != nil { + hasMoE = true + for _, e := range b.MoE.Routed.Experts { + wider(e) + } + for _, s := range b.MoE.Shared { + wider(s) + } + if nExperts == 0 { + nExperts, topK = len(b.MoE.Routed.Experts), b.MoE.Routed.TopK + } + } + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: cfg.Heads, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.mla = true + if hasMoE { + d.moe, d.moeRaw = true, true // DeepSeek-V2 raw-softmax·routedScale gating + d.nExperts, d.topK = nExperts, topK + d.routedScale = 1 + if cfg.RoutedScale > 0 { + d.routedScale = float32(cfg.RoutedScale) + } + } + d.qkNope, d.qkRope, d.vHead, d.qkHead = cfg.QKNope, cfg.QKRope, cfg.VHead, qkHead + d.qLoRA, d.kvLoRA = cfg.QLoraRank, cfg.KVLoraRank + scale := cfg.SoftmaxScale + if scale <= 0 { + scale = 1.0 / math.Sqrt(float64(qkHead)) + } + d.mlaScale = float32(scale) + // decoupled RoPE frequencies over QKRope (single-head convention, Heads:1) + invF64, posDiv64 := backend.RoPEFreqs(cfg.QKRope, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: 1}) + d.mlaPosDiv = float32(posDiv64) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + // MLA scratch + the QKRope frequency table. + c := d.maxLen + mlaInv := make([]float32, cfg.QKRope/2) + for i := range mlaInv { + mlaInv[i] = float32(invF64[i]) + } + d.mlaInv = mk(mlaInv) + d.mlaCQ = mk(make([]float32, c*cfg.QLoraRank)) + d.mlaQ = mk(make([]float32, c*cfg.Heads*qkHead)) + d.mlaComp = mk(make([]float32, c*(cfg.KVLoraRank+cfg.QKRope))) + d.mlaLatent = mk(make([]float32, c*cfg.KVLoraRank)) + d.mlaKV = mk(make([]float32, c*cfg.Heads*kvHead)) + d.mlaAttn = mk(make([]float32, c*cfg.Heads*cfg.VHead)) + + //perfscan:ignore PS3032 newDeepSeekV2Decoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqA: lin(b.WqA), wqB: lin(b.WqB), wkvA: lin(b.WkvA), wkvB: lin(b.WkvB), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, gFFN: mk(flat1D(b.PostAttnNorm.Gamma)).b, + gQA: mk(flat1D(b.QANorm.Gamma)).b, gKvA: mk(flat1D(b.KvANorm.Gamma)).b, + // MLA cache: K is heads·qkHead wide, V is heads·VHead wide (rectangular). + kC: mk(make([]float32, c*cfg.Heads*qkHead)).b, vC: mk(make([]float32, c*cfg.Heads*cfg.VHead)).b, + } + if b.Dense != nil { // dense SwiGLU FFN + gb.wG, gb.wU, gb.wD = lin(b.Dense.Wgate), lin(b.Dense.Wup), lin(b.Dense.Wdown) + } else { // DeepSeekMoE: routed top-k experts + ungated shared expert(s) + experts := make([]moeFFN, len(b.MoE.Routed.Experts)) + for i, ex := range b.MoE.Routed.Experts { + wgu, wg, wu := gateUp(ex.Wgate, ex.Wup) + experts[i] = moeFFN{wGU: wgu, wG: wg, wU: wu, wD: lin(ex.Wdown)} + } + gb.moeRouter, gb.moeExperts = d.f32Lin(mk)(b.MoE.Routed.Router.W), experts // router stays f32 (routing is selection-sensitive) + if len(b.MoE.Shared) > 0 { // realized as a single fused SwiGLU; ungated (moeSharedGate nil) + s := b.MoE.Shared[0] + gb.moeShared = moeFFN{wG: lin(s.Wgate), wU: lin(s.Wup), wD: lin(s.Wdown)} + } + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + d.out = lin(m.LmHead) + d.allocScratch(mk) // standard scratch (dx/xn/xn2/gate/up/mo/logits reused; Llama-shaped q/qkv unused) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newMambaDecoder builds a device Decoder for an nlp.Mamba — the FIRST non-transformer GPU decoder. +// Each layer is a pre-RMSNorm + selective-state-space mixer (nn.MambaBlock): no attention, no KV +// cache, no FFN. Decode is a linear-time recurrence, so the block records the per-timestep conv/SSM +// decode kernels (Conv1DStep/SSMStep, cuda-only) carrying per-block state across Step calls. Only +// DtProj and the conv carry a bias (the HF checkpoint is otherwise bias-free); A = −exp(A_log) is +// precomputed on the host. The LM head is tied (logits = hidden · Embedᵀ). cuda-only. +func newMambaDecoder(m *nlp.Mamba, ops backendOps) (*Decoder, error) { + cfg := m.Config + if len(m.Layers) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Mamba has no layers", ops.name) + } + mix0 := m.Layers[0].Mixer + // Mamba has no attention: Heads=1 satisfies newDecoderCommon's geometry (all attention/KV scratch + // is skipped by allocMambaScratch). Hidden is unused (no FFN). Ctx is a pos bound only — the + // recurrence has no cache, so allocMambaScratch never scales with it; pick a generous limit. + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: 1 << 20, Dim: cfg.DModel, Heads: 1, KVHeads: 1, + Layers: cfg.Layers, Hidden: 1, Eps: cfg.Eps, + } + d, derr := newDecoderCommon(lc, m.Embed, ops) + if derr != nil { + return nil, derr + } + d.mamba = true + d.dInner, d.mambaN, d.dConv, d.dtRank = mix0.DInner, mix0.N, mix0.DConv, mix0.DtRank + + var err error + mk := d.mkBuf(&err) + // lin builds a projection: the f32 default, or (Q8 mode) a resident Q8_0 GEMV. Mamba decode is + // weight-bandwidth-bound (~230 MB of f32 weights streamed per token), so quantizing the projections + // to Q8_0 cuts that ~4× — the dominant decode lever. Only the matmul weights go Q8; the tiny + // per-channel buffers (A, dt_bias, conv, dskip, norms) stay f32 (not GEMVs, negligible bandwidth). + lin := func(w *tensor.Tensor) linear { + in, out := w.Shape()[0], w.Shape()[1] + if ops.quantizeF32 != nil { + qw, qe := ops.quantizeF32(w) + if qe != nil { + if err == nil { + err = qe + } + return f32Linear{k: in, n: out} // err set → constructor bails before use + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } + //perfscan:ignore PS3032 newMambaDecoder block-upload, one-time model build + for _, ly := range m.Layers { + mx := ly.Mixer + // A = −exp(A_log), precomputed on the host (ALog is [dInner, N], row-major — the SSMStep layout). + alog := flat2D(mx.ALog) + A := make([]float32, len(alog)) + for i, v := range alog { + A[i] = float32(-math.Exp(float64(v))) + } + gb := block{ + gAttn: mk(flat1D(ly.Norm.Gamma)).b, // pre-mixer RMSNorm γ (stored in gAttn; recordMambaBlock norms dx→xn) + mInX: lin(mx.InX.W), mInZ: lin(mx.InZ.W), + mDtLow: lin(mx.DtLow.W), mDtProj: lin(mx.DtProj.W), + mBProj: lin(mx.BProj.W), mCProj: lin(mx.CProj.W), + mOutProj: lin(mx.OutProj.W), + mConvW: mk(flat2D(mx.ConvW)).b, mConvB: mk(flat1D(mx.ConvB)).b, + mA: mk(A).b, mDskip: mk(flat1D(mx.Dskip)).b, mDtBias: mk(flat1D(mx.DtProj.B)).b, + mConvState: mk(make([]float32, d.dInner*(d.dConv-1))).b, + mSsmState: mk(make([]float32, d.dInner*d.mambaN)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Head) // tied head [d_model, vocab]: logits = hidden · Embedᵀ + d.allocMambaScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newJambaDecoder builds a device Decoder for an nlp.Jamba — the hybrid stack that interleaves Mamba +// selective-scan mixers with NoPE grouped-query attention, and per layer runs either a sparse-MoE or a +// dense SwiGLU FFN. Each layer is pre-norm(input_layernorm) → mixer → pre-norm(pre_ff_layernorm) → FFN, +// residual-added. Because some layers are Mamba (state carried across tokens), the whole model decodes +// sequentially like Mamba (rows==1 Step; StepN loops Step). Attention layers keep a growing KV cache; +// Mamba layers keep conv/SSM state. The MoE router uses raw (non-renormalized) softmax weights. +func newJambaDecoder(m *nlp.Jamba, ops backendOps) (*Decoder, error) { + cfg := m.Config + if len(m.Layers) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Jamba has no layers", ops.name) + } + // Infer the shared dims from whichever layers carry them: a Mamba layer fixes the SSM geometry, + // any FFN fixes the hidden width, a MoE layer fixes the expert count. + var mix0 *nn.MambaBlock + var hidden, nExperts int + for _, ly := range m.Layers { + if ly.Mamba != nil && mix0 == nil { + mix0 = ly.Mamba.Block + } + if ly.Dense != nil && hidden == 0 { + hidden = ly.Dense.Wgate.Shape()[1] + } + if ly.MoE != nil { + if hidden == 0 { + hidden = ly.MoE.Experts[0].Wgate.Shape()[1] + } + if nExperts == 0 { + nExperts = len(ly.MoE.Experts) + } + } + } + if mix0 == nil { + return nil, fmt.Errorf("llamagpu(%s): Jamba has no Mamba layer", ops.name) + } + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + topK := cfg.TopK + if topK <= 0 { + topK = 2 + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: 1 << 16, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.jamba = true + d.noRope = true // NoPE attention; also skips the (unused) RoPE-frequency upload in allocScratch + d.dInner, d.mambaN, d.dConv, d.dtRank = mix0.DInner, mix0.N, mix0.DConv, mix0.DtRank + if nExperts > 0 { + d.moe, d.moeRaw = true, true // Jamba routes raw softmax weights (no top-k renormalization) + d.nExperts, d.topK, d.routedScale = nExperts, topK, 1 + } + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) // Q8_0 when ops.quantizeF32 set (NewJambaQ8CUDA), else f32 + //perfscan:ignore PS3032 newJambaDecoder block-upload, one-time model build + for _, ly := range m.Layers { + gb := block{ + gAttn: mk(flat1D(ly.InputNorm.Gamma)).b, // input_layernorm (mixer pre-norm) + gFFN: mk(flat1D(ly.PreFFNorm.Gamma)).b, // pre_ff_layernorm (FFN pre-norm) + } + if ly.Mamba != nil { // Mamba mixer: the Mamba projections + Jamba's dt/b/c RMSNorms + carried state + jm := ly.Mamba + mx := jm.Block + alog := flat2D(mx.ALog) + A := make([]float32, len(alog)) + for i, v := range alog { + A[i] = float32(-math.Exp(float64(v))) + } + gb.jMamba = true + gb.mInX, gb.mInZ = lin(mx.InX.W), lin(mx.InZ.W) + gb.mDtLow, gb.mDtProj = lin(mx.DtLow.W), lin(mx.DtProj.W) + gb.mBProj, gb.mCProj = lin(mx.BProj.W), lin(mx.CProj.W) + gb.mOutProj = lin(mx.OutProj.W) + gb.mConvW, gb.mConvB = mk(flat2D(mx.ConvW)).b, mk(flat1D(mx.ConvB)).b + gb.mA, gb.mDskip, gb.mDtBias = mk(A).b, mk(flat1D(mx.Dskip)).b, mk(flat1D(mx.DtProj.B)).b + gb.mConvState = mk(make([]float32, d.dInner*(d.dConv-1))).b + gb.mSsmState = mk(make([]float32, d.dInner*d.mambaN)).b + gb.mDtNorm, gb.mBNorm, gb.mCNorm = mk(flat1D(jm.DtNorm.Gamma)).b, mk(flat1D(jm.BNorm.Gamma)).b, mk(flat1D(jm.CNorm.Gamma)).b + } else { // NoPE grouped-query attention: bias-free q/k/v/o + a growing KV cache + at := ly.Attn + gb.wq, gb.wk, gb.wv, gb.wo = lin(at.Wq), lin(at.Wk), lin(at.Wv), lin(at.Wo) + gb.kC = mk(make([]float32, d.maxLen*d.kvDim)).b + gb.vC = mk(make([]float32, d.maxLen*d.kvDim)).b + } + if ly.MoE != nil { // sparse MoE FFN (raw-softmax top-k, no shared expert) + experts := make([]moeFFN, len(ly.MoE.Experts)) + for i, ex := range ly.MoE.Experts { + experts[i] = moeFFN{wG: lin(ex.Wgate), wU: lin(ex.Wup), wD: lin(ex.Wdown)} + } + gb.moeRouter, gb.moeExperts = d.f32Lin(mk)(ly.MoE.Router), experts // router stays f32 (routing is selection-sensitive) + } else { // dense SwiGLU FFN + gb.wG, gb.wU, gb.wD = lin(ly.Dense.Wgate), lin(ly.Dense.Wup), lin(ly.Dense.Wdown) + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Out) // untied lm_head [dim, vocab] + d.allocScratch(mk) // attention + FFN + (d.moe) MoE scratch, all maxLen-sized + // Mamba-mixer intermediates (rows==1): the only scratch allocScratch doesn't cover. + dI, N, dtR := d.dInner, d.mambaN, d.dtRank + d.mbXin, d.mbZ, d.mbXc = mk(make([]float32, dI)), mk(make([]float32, dI)), mk(make([]float32, dI)) + d.mbDtLow, d.mbDelta = mk(make([]float32, dtR)), mk(make([]float32, dI)) + d.mbB, d.mbC, d.mbY = mk(make([]float32, N)), mk(make([]float32, N)), mk(make([]float32, dI)) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newMamba2Decoder builds a device Decoder for an nlp.Mamba2 — the state-space-duality sibling of +// Mamba. Each layer is a pre-RMSNorm + SSD mixer: fused in_proj → [z | xBC | dt] → causal-conv + SiLU +// → [x | B | C] → Δ=softplus(dt+dt_bias) → the SSD scan (scalar per-head decay, B/C shared across a +// group, +D·x skip; via cu_ssd_step) → gated RMSNorm norm(y·SiLU(z)) → out_proj. No attention, no KV +// cache — decode is a linear-time recurrence, so it reuses the Mamba rows==1 sequential Step/StepN +// plumbing. in_proj/out_proj are torch [out,in] (transposed on upload); A=−exp(A_log); tied LM head. +func newMamba2Decoder(m *nlp.Mamba2, ops backendOps) (*Decoder, error) { + cfg := m.Config + if len(m.Layers) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Mamba2 has no layers", ops.name) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: 1 << 20, Dim: cfg.DModel, Heads: 1, KVHeads: 1, + Layers: cfg.Layers, Hidden: 1, Eps: cfg.Eps, + } + d, derr := newDecoderCommon(lc, m.Embed, ops) + if derr != nil { + return nil, derr + } + d.mamba2 = true + d.m2H, d.m2P, d.m2G, d.m2N = cfg.NumHeads, cfg.HeadDim, cfg.NGroups, cfg.N + d.m2Conv, d.m2Inter, d.m2CD = cfg.DConv, cfg.Intermediate, cfg.Intermediate+2*cfg.NGroups*cfg.N + + var err error + mk := d.mkBuf(&err) + // linT uploads a torch [out,in] weight transposed to the [in,out] the recorder matmul expects. + linT := func(w *tensor.Tensor) linear { + out, in := w.Shape()[0], w.Shape()[1] + if ops.quantizeF32 != nil { // Q8 mode: transpose the [out,in] weight to [in,out], then quantize (§ weight-bandwidth lever) + wT, te := w.Transpose(0, 1) // zero-copy view; NewResidentBQ8 materializes it via Contiguous() + if te != nil { + if err == nil { + err = te + } + return f32Linear{k: in, n: out} + } + qw, qe := ops.quantizeF32(wT) + if qe != nil { + if err == nil { + err = qe + } + return f32Linear{k: in, n: out} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } + return f32Linear{w: mk(flat2DT(w)).b, k: in, n: out} + } + //perfscan:ignore PS3032 newMamba2Decoder block-upload, one-time model build + for _, ly := range m.Layers { + mx := ly.Mixer + aLog := flat1D(mx.ALog) + A := make([]float32, len(aLog)) + for i, v := range aLog { + A[i] = float32(-math.Exp(float64(v))) // A = −exp(A_log), per head + } + gb := block{ + gAttn: mk(flat1D(ly.Norm.Gamma)).b, // pre-mixer RMSNorm γ + m2InProj: linT(mx.InProj), m2OutProj: linT(mx.OutProj), + m2ConvW: mk(flat2D(mx.ConvW)).b, m2ConvB: mk(flat1D(mx.ConvB)).b, + m2A: mk(A).b, m2DtBias: mk(flat1D(mx.DtBias)).b, m2Dskip: mk(flat1D(mx.D)).b, + m2NormW: mk(flat1D(mx.NormW)).b, + m2ConvState: mk(make([]float32, d.m2CD*(d.m2Conv-1))).b, + m2SsdState: mk(make([]float32, d.m2H*d.m2N*d.m2P)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + // tied head: m.Head is already [d_model, vocab] = Embedᵀ, my [in=d_model, out=vocab] orientation. + if ops.quantizeF32 != nil { + if qw, qe := ops.quantizeF32(m.Head); qe != nil { + if err == nil { + err = qe + } + } else { + d.qweights = append(d.qweights, qw) + d.out = quantLinear{w: qw} + } + } else { + d.out = f32Linear{w: mk(flat2D(m.Head)).b, k: cfg.DModel, n: cfg.Vocab} + } + d.allocMamba2Scratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newRWKVDecoder builds a device Decoder for an nlp.RWKV — GoAI's third recurrent family (after the +// Mamba SSMs). Each layer is an nn.RWKVBlock: a WKV time-mixing sublayer (token-shift → r/k/v → the +// stabilized WKV recurrence via cu_wkv_step → receptance gate → output proj) and a gated squared-ReLU +// channel-mixing sublayer, each with its own LayerNorm + residual. No attention, no KV cache, no +// positional encoding — decode is O(1) recurrent, reusing the rows==1 sequential Step/StepN plumbing. +// An extra ln0 LayerNorm transforms the embedding before block 0; the head is UNTIED. cuda-only. +func newRWKVDecoder(m *nlp.RWKV, ops backendOps) (*Decoder, error) { + cfg := m.Config + if len(m.Blocks) == 0 { + return nil, fmt.Errorf("llamagpu(%s): RWKV has no blocks", ops.name) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: 1 << 20, Dim: cfg.Dim, Heads: 1, KVHeads: 1, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, + } + d, derr := newDecoderCommon(lc, m.Embed, ops) + if derr != nil { + return nil, derr + } + d.rwkv = true + d.lnBias = true // final norm (ln_out) is a full LayerNorm + d.rwHidden = cfg.Hidden + + var err error + mk := d.mkBuf(&err) + lin := func(w *tensor.Tensor) linear { // nn.Linear.W is already [in,out] (transposed at HF load) + in, out := w.Shape()[0], w.Shape()[1] + if ops.quantizeF32 != nil { // Q8 mode: same weight-bandwidth lever as Mamba (time-mix/channel-mix GEMVs) + qw, qe := ops.quantizeF32(w) + if qe != nil { + if err == nil { + err = qe + } + return f32Linear{k: in, n: out} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } + // mu uploads the interpolator and its 1−μ complement (token-shift compose needs both). + comp := func(mu *tensor.Tensor) buffer { + v := flat1D(mu) + o := make([]float32, len(v)) + for i, x := range v { + o[i] = 1 - x + } + return mk(o).b + } + expv := func(t *tensor.Tensor) buffer { // per-channel decay w = exp(WLog) + v := flat1D(t) + o := make([]float32, len(v)) + for i, x := range v { + o[i] = float32(math.Exp(float64(x))) + } + return mk(o).b + } + zeros := func() buffer { return mk(make([]float32, d.d)).b } + //perfscan:ignore PS3032 newRWKVDecoder block-upload, one-time model build + for _, bl := range m.Blocks { + gb := block{ + gAttn: mk(flat1D(bl.LN1.Gamma)).b, bAttn: mk(flat1D(bl.LN1.Beta)).b, // LN1 + gFFN: mk(flat1D(bl.LN2.Gamma)).b, bFFN: mk(flat1D(bl.LN2.Beta)).b, // LN2 + wq: lin(bl.Wr.W), wk: lin(bl.Wk.W), wv: lin(bl.Wv.W), wo: lin(bl.Wo.W), // time-mix Wr/Wk/Wv/Wo + wG: lin(bl.CWk.W), wD: lin(bl.CWv.W), rwCWr: lin(bl.CWr.W), // channel-mix CWk/CWv/CWr + rwMuR: mk(flat1D(bl.MuR)).b, rwMuK: mk(flat1D(bl.MuK)).b, rwMuV: mk(flat1D(bl.MuV)).b, + rwCMuR: mk(flat1D(bl.CMuR)).b, rwCMuK: mk(flat1D(bl.CMuK)).b, + rwOmR: comp(bl.MuR), rwOmK: comp(bl.MuK), rwOmV: comp(bl.MuV), + rwOmCR: comp(bl.CMuR), rwOmCK: comp(bl.CMuK), + rwW: expv(bl.WLog), rwU: mk(flat1D(bl.U)).b, + rwPrevTM: zeros(), rwPrevCM: zeros(), rwAA: zeros(), rwBB: zeros(), rwPP: zeros(), + } + d.blocks = append(d.blocks, gb) + } + d.rwPreLNg = mk(flat1D(m.PreLN.Gamma)) // ln0 (before block 0) + d.rwPreLNb = mk(flat1D(m.PreLN.Beta)) + d.gFinal = mk(flat1D(m.LNOut.Gamma)) // final LayerNorm ln_out + d.bFinal = mk(flat1D(m.LNOut.Beta)) + d.out = lin(m.Head) // untied head [dim, vocab] + d.allocRWKVScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newFalconDecoder builds a device Decoder for an nlp.Falcon (falcon-7b class: parallel_attn=True, +// multi_query=True, new_decoder_architecture=False). Every departure is an existing generalization: +// SINGLE-NORM parallel residual (parallelRes — one input_layernorm feeds attention AND the MLP, both +// summed onto x0, like Cohere), full LayerNorm WITH bias (lnBias), a 2-layer GELU MLP (ffnGELU), and +// multi-query attention (MQA = GQA with one KV head, kvH=1, from the fused query_key_value split). +// Full split-half rope, no linear biases, untied lm_head. cuda-only (parallel residual + LayerNorm). +func newFalconDecoder(m *nlp.Falcon, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = 1 // Falcon MQA + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // LayerNorm with bias + d.ffnGELU = true // 2-layer GELU MLP + d.parallelRes = true // one-norm parallel residual: input_layernorm feeds attn AND MLP + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + //perfscan:ignore PS3032 newFalconDecoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + // gFFN/bFFN unused: parallel one-norm reuses the attn norm (d.xn) for the MLP. + wG: lin(b.Wh), wD: lin(b.Wout), // 2-layer GELU MLP: wG = dense_h_to_4h, wD = dense_4h_to_h + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + d.bFinal = mk(flat1D(m.FinalNorm.Beta)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newOLMo2Decoder builds a device Decoder for an nlp.OLMo2 (Allen AI). OLMo2 shares Llama's SwiGLU/ +// GQA/RoPE core but reshuffles the residual structure in two ways, both new gated generalizations: +// - POST-NORM (postNorm): there is no input_layernorm; each sublayer reads the RAW residual and +// its OUTPUT is normed (post_attention_layernorm / post_feedforward_layernorm) before the add. +// - FULL-WIDTH QK-norm (qkNorm+qkNormFull): an RMSNorm over the ENTIRE q_proj / k_proj output +// (not per-head like Qwen3), applied before RoPE. +// +// The lm_head is untied. cuda-only. +func newOLMo2Decoder(m *nlp.OLMo2, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.postNorm = true // sublayer output is normed before the residual add + d.qkNorm = true // RMSNorm on the q/k projections before RoPE... + d.qkNormFull = true // ...over the full width, not per head + // decoupled head_dim (read back from q_proj; harmless when head_dim == dim/heads) + if qOut := m.Blocks[0].Wq.Shape()[1]; qOut%cfg.Heads == 0 && qOut/cfg.Heads != d.dk { + hd := qOut / cfg.Heads + d.dk, d.qDim, d.kvDim, d.half, d.rotaryDim = hd, d.h*hd, d.kvH*hd, hd/2, hd + d.scale = float32(1.0 / math.Sqrt(float64(hd))) + invF64, posDiv64 := backend.RoPEFreqs(hd, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, hd/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + } + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + //perfscan:ignore PS3032 newOLMo2Decoder block-upload, one-time model build + for _, b := range m.Blocks { + wgu, wg, wu := gateUp(b.FFN.Wgate, b.FFN.Wup) + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + // post-norm: gAttn = post_attention_layernorm (on the attn output), gFFN = post_feedforward_layernorm. + gAttn: mk(flat1D(b.PostAttnNorm.Gamma)).b, gFFN: mk(flat1D(b.PostFFNNorm.Gamma)).b, + qN: mk(flat1D(b.QNorm.Gamma)).b, kN: mk(flat1D(b.KNorm.Gamma)).b, // full-width [qDim]/[kvDim] + wGU: wgu, wG: wg, wU: wu, wD: lin(b.FFN.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newGemma2Decoder builds a device Decoder for an nlp.Gemma2. Gemma2 extends Gemma with three more +// departures, all reusing the generalizations added for its siblings: +// - SANDWICH norms (sandwich): each sublayer is pre-normed AND its output post-normed — four +// (1+w) RMSNorms per block (input/post_attention, pre_feedforward/post_feedforward). +// - attention-logit SOFT-CAP (attnCap → MHACap): scaled scores pass through cap·tanh(·/cap) +// before the mask+softmax. +// - query_pre_attn_scalar: the pre-softmax scale is 1/√scalar (not 1/√head_dim in general). +// +// Plus Gemma's √dim embedding normalizer, (1+w) RMSNorm, GeGLU, decoupled head_dim and tied lm_head. +// The final-logit soft-cap is a MONOTONIC map, so it never changes the greedy argmax and is omitted +// (greedy generation is identical with or without it). cuda-only (the soft-cap kernel is cuda-only). +func newGemma2Decoder(m *nlp.Gemma2, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.HeadDim + if headDim <= 0 { + headDim = cfg.Dim / cfg.Heads + } + if qOut := m.Blocks[0].Wq.Shape()[1]; qOut%cfg.Heads == 0 { + headDim = qOut / cfg.Heads + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.FFN, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.ffnGEGLU = true + d.sandwich = true + d.embMult = float32(math.Sqrt(float64(cfg.Dim))) + if cfg.AttnLogitCap > 0 { + d.attnCap = float32(cfg.AttnLogitCap) + } + // decoupled head_dim geometry override (as Gemma) + if headDim != d.dk { + d.dk, d.qDim, d.kvDim, d.half, d.rotaryDim = headDim, d.h*headDim, d.kvH*headDim, headDim/2, headDim + invF64, posDiv64 := backend.RoPEFreqs(headDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, headDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + } + // query_pre_attn_scalar sets the pre-softmax scale to 1/√scalar (0 → head_dim = standard). + qpa := cfg.QueryPreAttnScalar + if qpa == 0 { + qpa = float64(headDim) + } + d.scale = float32(1.0 / math.Sqrt(qpa)) + + var err error + mk := d.mkBuf(&err) + fused := d.mkFused(mk, ops, &err) + lin := d.mkLin(mk, ops, &err) + //perfscan:ignore PS3032 newGemma2Decoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, gAttn2: mk(flat1D(b.PostAttnNorm.Gamma)).b, + gFFN: mk(flat1D(b.PreFFNNorm.Gamma)).b, gFFN2: mk(flat1D(b.PostFFNNorm.Gamma)).b, + wG: lin(b.FFN.Wgate), wU: lin(b.FFN.Wup), wD: lin(b.FFN.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + // tied lm_head: Out = TokEmbᵀ [dim, vocab] + vocab, dim := m.TokEmb.Shape()[0], m.TokEmb.Shape()[1] + // Route the tied lm_head through the same quantize-aware builder so the Q8 constructors quantize it too + // — Gemma's vocab is huge (256k), so an f32 lm_head would dominate decode bandwidth and dilute Q8. The + // f32 path is byte-identical (mkLin with no hook == the old f32Linear{mk(outW)}). + outT := tensor.New(tensor.F32, tensor.Shape{dim, vocab}) + flat2DTInto(outT.Storage().F32()[:dim*vocab], m.TokEmb) + d.out = lin(outT) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newMPTDecoder builds a device Decoder for an nlp.MPT (MosaicML). MPT is the first ALiBi decoder: +// position enters ONLY through a per-head linear bias on the attention scores (no RoPE, no positional +// embedding), so newMPTDecoder sets noRope and uploads the backend.ALiBiSlopes so recordMHA routes +// attention through MHAALiBi. Otherwise it is a sequential-residual block with weight-only LayerNorm +// (lnBias, β=0), standard MHA (no GQA), a bias-free 2-layer GELU MLP (ffnGELU) and a tied lm_head. +// cuda-only (the ALiBi attention kernel is cuda-only). +func newMPTDecoder(m *nlp.MPT, ops backendOps) (*Decoder, error) { + cfg := m.Config + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: cfg.Heads, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: 10000, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // weight-only LayerNorm (β = 0) + d.ffnGELU = true // bias-free 2-layer GELU MLP + d.noRope = true // ALiBi, no rotary + + var err error + mk := d.mkBuf(&err) + slopes := make([]float32, d.h) + for i, s := range backend.ALiBiSlopes(d.h) { + slopes[i] = float32(s) + } + d.aliBiSlopes = mk(slopes).b + + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + //perfscan:ignore PS3032 newMPTDecoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.Norm1.Gamma)).b, bAttn: mk(flat1D(b.Norm1.Beta)).b, + gFFN: mk(flat1D(b.Norm2.Gamma)).b, bFFN: mk(flat1D(b.Norm2.Beta)).b, + wG: lin(b.Wup), wD: lin(b.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + d.bFinal = mk(flat1D(m.FinalNorm.Beta)) + d.out = lin(m.Out) // tied wteᵀ, provided [dim, vocab] by the loader + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newMixtralDecoder builds a device Decoder for an nlp.Mixtral (sparse Mixture-of-Experts). The +// attention/norm stack is the plain Llama core (RMSNorm, fused QKV, full rope, GQA, and the optional +// per-head QK-norm shared with Qwen3-MoE); only the FFN sublayer becomes a sparse MoE: a router +// scores the experts, MoEGate picks the renormalized top-k weights, and every expert's SwiGLU is +// weighted-combined onto the residual (recordMoE). Dense expert eval — correct but not yet the sparse +// top-k gather. cuda-only (the MoE routing/combine kernels are cuda-only). +func newMixtralDecoder(m *nlp.Mixtral, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + if len(m.Blocks) == 0 || m.Blocks[0].MoE == nil || len(m.Blocks[0].MoE.Experts) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Mixtral has no MoE experts", ops.name) + } + // Derive the expert count / top-k from the loaded SparseMoE (MixtralFromHF infers Experts and + // may not write it back into Config). + nExperts := len(m.Blocks[0].MoE.Experts) + topK := m.Blocks[0].MoE.TopK + if topK <= 0 { + topK = 2 + } + hidden := cfg.Hidden + if hidden <= 0 { // inferred, may not be written back to Config — read it from an expert's Wgate [dim, ffn] + hidden = m.Blocks[0].MoE.Experts[0].Wgate.Shape()[1] + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.moe = true + d.nExperts = nExperts + d.topK = topK + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + qkGain := func(n *nn.RMSNorm) buffer { // optional per-head QK-norm (Qwen3-MoE); nil for Mixtral + if n == nil { + return nil + } + d.qkNorm = true + return mk(flat1D(n.Gamma)).b + } + //perfscan:ignore PS3032 newMixtralDecoder block-upload, one-time model build + for _, b := range m.Blocks { + experts := make([]moeFFN, len(b.MoE.Experts)) + for i, ex := range b.MoE.Experts { + wgu, wg, wu := gateUp(ex.Wgate, ex.Wup) + experts[i] = moeFFN{wGU: wgu, wG: wg, wU: wu, wD: lin(ex.Wdown)} + } + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + qN: qkGain(b.QNorm), kN: qkGain(b.KNorm), + moeRouter: d.f32Lin(mk)(b.MoE.Router.W), moeExperts: experts, // router stays f32 (routing is selection-sensitive) + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newQwen2MoEDecoder builds a device Decoder for an nlp.Qwen2MoE. Qwen2-MoE is a routed sparse MoE +// PLUS a shared expert: the FFN output is sparse_moe(x) + sigmoid(x·gate)·shared_expert(x), a single +// SwiGLU run on every token and scaled by a per-token sigmoid gate (recordMoE's shared-expert tail). +// Attention is the Llama core with Qwen2's q/k/v projection biases (o_proj bias-free). cuda-only. +func newQwen2MoEDecoder(m *nlp.Qwen2MoE, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + if len(m.Blocks) == 0 || m.Blocks[0].MoE == nil || len(m.Blocks[0].MoE.Experts) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Qwen2MoE has no MoE experts", ops.name) + } + b0 := m.Blocks[0] + nExperts := len(b0.MoE.Experts) + topK := b0.MoE.TopK + if topK <= 0 { + topK = 2 + } + hidden := b0.MoE.Experts[0].Wgate.Shape()[1] + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.moe = true + d.nExperts = nExperts + d.topK = topK + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + fusedBias := func(bq, bk, bv *tensor.Tensor) buffer { + if bq == nil && bk == nil && bv == nil { + return nil + } + return mk(append(append(append([]float32{}, flat1D(bq)...), flat1D(bk)...), flat1D(bv)...)).b + } + //perfscan:ignore PS3032 newQwen2MoEDecoder block-upload, one-time model build + for _, b := range m.Blocks { + experts := make([]moeFFN, len(b.MoE.Experts)) + for i, ex := range b.MoE.Experts { + wgu, wg, wu := gateUp(ex.Wgate, ex.Wup) + experts[i] = moeFFN{wGU: wgu, wG: wg, wU: wu, wD: lin(ex.Wdown)} + } + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), qkvBias: fusedBias(b.Bq, b.Bk, b.Bv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + moeRouter: d.f32Lin(mk)(b.MoE.Router.W), moeExperts: experts, // router stays f32 (routing is selection-sensitive) + moeShared: moeFFN{wG: lin(b.Shared.Wgate), wU: lin(b.Shared.Wup), wD: lin(b.Shared.Wdown)}, + moeSharedGate: lin(b.SharedGate), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newGraniteMoEDecoder builds a device Decoder for an nlp.GraniteMoE — the MoE sibling of dense +// Granite: a plain Llama attention core + a sparse Mixture-of-Experts FFN, wrapped in the four +// Granite config scalars (all identity when 0/1, so this is newMixtralDecoder + the load-time folds +// from newDecoder's Granite handling). AttentionMult overrides the softmax scale; EmbeddingMult +// scales the gathered embedding; ResidualMult is baked into Wo AND every expert's Wdown (their +// matmul feeds the residual); LogitsScale is baked as 1/scale into the untied lm_head. cuda-only. +func newGraniteMoEDecoder(m *nlp.GraniteMoE, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + if len(m.Blocks) == 0 || m.Blocks[0].MoE == nil || len(m.Blocks[0].MoE.Experts) == 0 { + return nil, fmt.Errorf("llamagpu(%s): GraniteMoE has no MoE experts", ops.name) + } + b0 := m.Blocks[0] + nExperts := len(b0.MoE.Experts) + topK := b0.MoE.TopK + if topK <= 0 { + topK = 2 + } + hidden := b0.MoE.Experts[0].Wgate.Shape()[1] + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.moe = true + d.nExperts = nExperts + d.topK = topK + if cfg.AttentionMult != 0 { + d.scale = float32(cfg.AttentionMult) + } + if cfg.EmbeddingMult != 0 && cfg.EmbeddingMult != 1 { + d.embMult = float32(cfg.EmbeddingMult) + } + resMult := float32(1) + if cfg.ResidualMult != 0 && cfg.ResidualMult != 1 { + resMult = float32(cfg.ResidualMult) + } + outScale := float32(1) + if cfg.LogitsScale != 0 && cfg.LogitsScale != 1 { + outScale = float32(1 / cfg.LogitsScale) + } + + var err error + mk := d.mkBuf(&err) + gateUp := d.gateUpBuilder(mk, ops, &err) + linS := d.mkLinS(mk, ops, &err) // ResidualMult / logit-scale fold — quantize-aware (Q8 under the hook) + fused := d.mkFused(mk, ops, &err) + //perfscan:ignore PS3032 newGraniteMoEDecoder block-upload, one-time model build + for _, b := range m.Blocks { + experts := make([]moeFFN, len(b.MoE.Experts)) + for i, ex := range b.MoE.Experts { + wgu, wg, wu := gateUp(ex.Wgate, ex.Wup) + experts[i] = moeFFN{wGU: wgu, wG: wg, wU: wu, wD: linS(ex.Wdown, resMult)} + } + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: linS(b.Wo, resMult), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + moeRouter: d.f32Lin(mk)(b.MoE.Router.W), moeExperts: experts, // router stays f32 (routing is selection-sensitive) + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = linS(m.Out, outScale) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newOLMoEDecoder builds a device Decoder for an nlp.OLMoE (Allen AI sparse-MoE). OLMoE is the +// pre-norm sparse-MoE composition: the plain Llama attention core with FULL-WIDTH q/k RMSNorm +// (qkNorm+qkNormFull, one RMSNorm over the whole q/k projection — same as OLMo2) plus a sparse +// Mixture-of-Experts FFN (recordMoE) and an untied lm_head. cuda-only (MoE + full-width QK-norm). +func newOLMoEDecoder(m *nlp.OLMoE, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + if len(m.Blocks) == 0 || m.Blocks[0].MoE == nil || len(m.Blocks[0].MoE.Experts) == 0 { + return nil, fmt.Errorf("llamagpu(%s): OLMoE has no MoE experts", ops.name) + } + b0 := m.Blocks[0] + nExperts := len(b0.MoE.Experts) + topK := b0.MoE.TopK + if topK <= 0 { + topK = 2 + } + hidden := cfg.Hidden + if hidden <= 0 { + hidden = b0.MoE.Experts[0].Wgate.Shape()[1] + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.moe = true + d.nExperts = nExperts + d.topK = topK + d.qkNorm = true // full-width q/k RMSNorm before RoPE... + d.qkNormFull = true // ...over the whole projection (OLMo2/OLMoE) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + //perfscan:ignore PS3032 newOLMoEDecoder block-upload, one-time model build + for _, b := range m.Blocks { + experts := make([]moeFFN, len(b.MoE.Experts)) + for i, ex := range b.MoE.Experts { + wgu, wg, wu := gateUp(ex.Wgate, ex.Wup) + experts[i] = moeFFN{wGU: wgu, wG: wg, wU: wu, wD: lin(ex.Wdown)} + } + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + qN: mk(flat1D(b.QNorm.Gamma)).b, kN: mk(flat1D(b.KNorm.Gamma)).b, // full-width [qDim]/[kvDim] + moeRouter: d.f32Lin(mk)(b.MoE.Router.W), moeExperts: experts, // router stays f32 (routing is selection-sensitive) + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newStarCoder2Decoder builds a device Decoder for an nlp.StarCoder2: LayerNorm-with-bias + +// biased q/k/v/o projections + a biased 2-layer GELU MLP (c_fc → GELU → c_proj) + FULL rope + +// GQA + untied head. Every one of those is a core generalization (lnBias, biased projections, +// ffnGELU); rotary stays full (rotaryDim == dk). cuda-only. +func newStarCoder2Decoder(m *nlp.StarCoder2, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // LayerNorm-with-bias + d.ffnGELU = true // 2-layer GELU MLP + // rotary stays full (rotaryDim == dk, the newDecoderCommon default) — no inv override. + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + fusedBias := func(bq, bk, bv *tensor.Tensor) buffer { // [Bq | Bk | Bv] = [D+2·kvDim] + fb := append(append(append([]float32{}, flat1D(bq)...), flat1D(bk)...), flat1D(bv)...) + return mk(fb).b + } + //perfscan:ignore PS3032 newStarCoder2Decoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), qkvBias: fusedBias(b.Bq, b.Bk, b.Bv), + wo: lin(b.Wo), oBias: mk(flat1D(b.Bo)).b, + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + gFFN: mk(flat1D(b.PostAttnNorm.Gamma)).b, bFFN: mk(flat1D(b.PostAttnNorm.Beta)).b, + wG: lin(b.Wfc), fcBias: mk(flat1D(b.Bfc)).b, + wD: lin(b.Wproj), projBias: mk(flat1D(b.Bproj)).b, + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.bFinal = mk(flat1D(m.Norm.Beta)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newPhiDecoder builds a device Decoder for an nlp.Phi: the hardest new-arch shape so far — ONE +// input LayerNorm feeding BOTH attention and the MLP (one-norm parallel residual), biased q/k/v/dense +// projections, a biased 2-layer GELU MLP (fc1 → GELU → fc2), PARTIAL rotary, MHA, a biased untied +// lm_head and a final LayerNorm. Every one of those is a core generalization (parallelRes, biased +// projections, ffnGELU, lnBias, rotaryDim, outBias). cuda-only. +func newPhiDecoder(m *nlp.Phi, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.HeadDim + if headDim <= 0 { + headDim = cfg.Dim / cfg.Heads + } + rotaryDim := cfg.RotaryDim + if rotaryDim <= 0 { + pct := cfg.RotaryPct + if pct == 0 { + pct = 0.25 + } + rotaryDim = int(float64(headDim) * pct) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // LayerNorm-with-bias + d.ffnGELU = true // 2-layer GELU MLP + d.parallelRes = true // one-norm parallel residual (input_layernorm feeds attn AND mlp) + d.rotaryDim = rotaryDim + if d.rotaryDim <= 0 || d.rotaryDim > d.dk || d.rotaryDim%2 != 0 { + return nil, fmt.Errorf("llamagpu(%s): Phi rotaryDim %d invalid for headDim %d", ops.name, d.rotaryDim, d.dk) + } + invF64, posDiv64 := backend.RoPEFreqs(d.rotaryDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, d.rotaryDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + fusedBias := func(bq, bk, bv *tensor.Tensor) buffer { + fb := append(append(append([]float32{}, flat1D(bq)...), flat1D(bk)...), flat1D(bv)...) + return mk(fb).b + } + //perfscan:ignore PS3032 newPhiDecoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), qkvBias: fusedBias(b.Bq, b.Bk, b.Bv), + wo: lin(b.Wdense), oBias: mk(flat1D(b.Bdense)).b, + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + // gFFN/bFFN unused: parallel one-norm reuses the attn norm (d.xn) for the MLP. + wG: lin(b.Wfc1), fcBias: mk(flat1D(b.Bfc1)).b, + wD: lin(b.Wfc2), projBias: mk(flat1D(b.Bfc2)).b, + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + d.bFinal = mk(flat1D(m.FinalNorm.Beta)) + d.out = lin(m.Out) + d.outBias = mk(flat1D(m.OutBias)) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newGPTNeoXDecoder uploads an nlp.GPTNeoX onto the batched Decoder core. GPT-NeoX is the TWO-norm +// parallel residual: a separate input_layernorm feeds attention and post_attention_layernorm feeds +// the MLP, but BOTH read the raw residual x0 and both outputs sum onto it. Beyond that it is +// LayerNorm-with-bias, biased q/k/v/dense, a biased GELU MLP, (partial or full) rotary and an +// untied lm_head with no bias. The fourth new-arch GPU graph decoder — and the one that exercises +// parallelTwoNorm (the FFN norm must be taken from x0 before the attention output is added back). +func newGPTNeoXDecoder(m *nlp.GPTNeoX, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.Dim / cfg.Heads + rotaryDim := cfg.RotaryDim + if rotaryDim <= 0 { + pct := cfg.RotaryPct + if pct == 0 { + pct = 1 // GPT-NeoX default is full rope + } + rotaryDim = int(float64(headDim) * pct) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // LayerNorm-with-bias + d.ffnGELU = true // 2-layer GELU MLP + d.parallelTwoNorm = true // two-norm parallel residual (input_ln → attn, post_attn_ln → mlp, both over x0) + d.rotaryDim = rotaryDim + if d.rotaryDim <= 0 || d.rotaryDim > d.dk || d.rotaryDim%2 != 0 { + return nil, fmt.Errorf("llamagpu(%s): GPT-NeoX rotaryDim %d invalid for headDim %d", ops.name, d.rotaryDim, d.dk) + } + invF64, posDiv64 := backend.RoPEFreqs(d.rotaryDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, d.rotaryDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + fusedBias := func(bq, bk, bv *tensor.Tensor) buffer { + fb := append(append(append([]float32{}, flat1D(bq)...), flat1D(bk)...), flat1D(bv)...) + return mk(fb).b + } + //perfscan:ignore PS3032 newGPTNeoXDecoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), qkvBias: fusedBias(b.Bq, b.Bk, b.Bv), + wo: lin(b.Wo), oBias: mk(flat1D(b.Bo)).b, + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + gFFN: mk(flat1D(b.PostAttnNorm.Gamma)).b, bFFN: mk(flat1D(b.PostAttnNorm.Beta)).b, + wG: lin(b.Wh), fcBias: mk(flat1D(b.Bh)).b, + wD: lin(b.Wout), projBias: mk(flat1D(b.Bout)).b, + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + d.bFinal = mk(flat1D(m.FinalNorm.Beta)) + d.out = lin(m.Out) // untied embed_out, no bias + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newQuantDecoder uploads a quantized Llama: RMSNorm gains + device-resident KV caches, every +// projection as a RESIDENT quantized weight consumed by the record-mode QMatMulResident (§T415). +// Cache storage stays f32 unless the constructor supplies newF16KVBuffer; that opt-in is deliberately +// limited to dk=64, the Metal kernel shape proven by P-01M093C3MFF0E8E1THCWYA1757. +func newQuantDecoder(m *nlp.QuantLlama, ops backendOps) (*Decoder, error) { + if ops.uploadQWeight == nil { + return nil, fmt.Errorf("llamagpu(%s): backend has no resident quantized upload", ops.name) + } + cfg := m.Config + d, derr := newDecoderCommon(cfg, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + if ops.newF16KVBuffer != nil { + if d.dk != 64 { + return nil, fmt.Errorf("llamagpu(%s): f16 KV cache requires head dimension 64, got %d", ops.name, d.dk) + } + d.f16KV = true + } + var err error + mk := d.mkBuf(&err) + mkCache := func(n int) buffer { + if !d.f16KV { + return mk(make([]float32, n)).b + } + if err != nil { + return nil + } + b, e := ops.newF16KVBuffer(n) + if e != nil { + err = e + return nil + } + d.all = append(d.all, b) + return b + } + qlin := func(q *nn.QuantLinear) linear { // resident quantized [Out,In] weight + if err != nil { + return quantLinear{} + } + w, e := ops.uploadQWeight(q.Weight, uint32(q.QT), q.Out, q.In) + if e != nil { + err = e + return quantLinear{} + } + d.qweights = append(d.qweights, w) + return quantLinear{w: w} + } + // fused QKV for quantized blocks (§T613): ggml resident weights are [Out,In] ROW-major, so + // fusing = appending the raw quantized bytes (out rows q‖k‖v) — valid only when all three + // projections share one quant type; mixed-type blocks keep the unfused three-matmul chain. + qfused := func(q1, q2, q3 *nn.QuantLinear) linear { + if err != nil || q1.QT != q2.QT || q2.QT != q3.QT { + return nil + } + wb := make([]byte, 0, len(q1.Weight)+len(q2.Weight)+len(q3.Weight)) + wb = append(append(append(wb, q1.Weight...), q2.Weight...), q3.Weight...) + w, e := ops.uploadQWeight(wb, uint32(q1.QT), q1.Out+q2.Out+q3.Out, q1.In) + if e != nil { + err = e + return nil + } + d.qweights = append(d.qweights, w) + return quantLinear{w: w} + } + qgroup := func(weights ...linear) linear { + if err != nil || ops.groupQWeights == nil || len(weights) != 3 { + return nil + } + qw := make([]qweight, 3) + for i, w := range weights { + ql, ok := w.(quantLinear) + if !ok || ql.w == nil { + return nil + } + qw[i] = ql.w + } + g, e := ops.groupQWeights(qw) + if e == backend.ErrQuantUnsupported { + return nil + } + if e != nil { + err = e + return nil + } + d.qweights = append(d.qweights, g) + return quantLinear{w: g} + } + // Column-fused ffn_gate|ffn_up for quantized blocks (same trick as qfused: ggml resident weights are + // [Out,In] row-major, so fusing = appending the raw quant bytes of gate‖up out-rows). Only when the + // backend supports SwiGLUHalves (ops.fusedGateUp) and both share one quant type; else nil → the + // unfused two-GEMV wG/wU path. Mirrors newDecoder's wGU so the raw-quant decode matches its speed. + qfused2 := func(q1, q2 *nn.QuantLinear) linear { + if err != nil || !ops.fusedGateUp || q1.QT != q2.QT { + return nil + } + wb := make([]byte, 0, len(q1.Weight)+len(q2.Weight)) + wb = append(append(wb, q1.Weight...), q2.Weight...) + w, e := ops.uploadQWeight(wb, uint32(q1.QT), q1.Out+q2.Out, q1.In) + if e != nil { + err = e + return nil + } + d.qweights = append(d.qweights, w) + return quantLinear{w: w} + } + //perfscan:ignore PS3032 newQuantDecoder block-upload, one-time model build + for _, b := range m.Blocks { + wgu := qfused2(b.FFN.Gate, b.FFN.Up) + gb := block{ + wqkv: qfused(b.Wq, b.Wk, b.Wv), wo: qlin(b.Wo), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + wGU: wgu, wD: qlin(b.FFN.Down), + kC: mkCache(d.maxLen * d.kvDim), vC: mkCache(d.maxLen * d.kvDim), + } + if wgu == nil { // no SwiGLUHalves backend or mixed types: keep the two-GEMV gate/up path + gb.wG, gb.wU = qlin(b.FFN.Gate), qlin(b.FFN.Up) + } + if gb.wqkv == nil { // mixed quant types: keep exact raw projections for single-token decode + gb.wq, gb.wk, gb.wv = qlin(b.Wq), qlin(b.Wk), qlin(b.Wv) + gb.wqkvPrefill = qgroup(gb.wq, gb.wk, gb.wv) + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = qlin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// encodeStep records the whole single-token layer stack + vocab head for absolute position +// pos into a fresh command buffer WITHOUT committing it. The recorded commands depend only +// on pos — not on the token, whose embedding is read from d.dx at execution time — which is +// what makes the §T614 encode-overlap legal: while the GPU executes step pos, the host +// pre-encodes step pos+1, and the next Step only uploads dx and commits. +func (d *Decoder) encodeStep(pos int) (recorder, error) { + newRecorder := d.ops.newRecorder + if d.ops.newDecodeRecorder != nil { + newRecorder = d.ops.newDecodeRecorder + } + r, err := newRecorder() + if err != nil { + return nil, err + } + D, H, KVH, dk, kvDim := d.qDim, d.h, d.kvH, d.dk, d.kvDim // D = attention Q/O width (h·dk); = d.d unless head_dim decoupled + if d.rwkv { // RWKV: an extra LayerNorm (ln0) transforms the embedding once before block 0 + if e := r.LayerNorm(d.dx.b, d.rwPreLNg.b, d.rwPreLNb.b, d.dx.b, 1, d.d, d.eps); e != nil { + r.Free() + return nil, e + } + } + for _, b := range d.blocks { + if d.rwkv { // RWKV-4 block: WKV time-mix + gated squared-ReLU channel-mix, no attention/KV cache + if e := d.recordRWKVBlock(r, b); e != nil { + r.Free() + return nil, e + } + continue + } + if d.mamba { // Mamba selective-scan block: no attention/FFN split — the mixer is the whole layer + if e := d.recordMambaBlock(r, b); e != nil { + r.Free() + return nil, e + } + continue + } + if d.mamba2 { // Mamba-2 SSD block (scalar-decay, grouped B/C, gated RMSNorm) + if e := d.recordMamba2Block(r, b); e != nil { + r.Free() + return nil, e + } + continue + } + if d.jamba { // hybrid: per-block mixer (Mamba OR NoPE attention), then a per-block MoE/dense FFN + var e error + if b.jMamba { + e = d.recordMambaBlock(r, b) // Mamba mixer w/ Jamba dt/b/c norms → dx += mixer(inputNorm(dx)) + } else { + e = d.recordJambaAttn(r, b, pos) // NoPE GQA attention → dx += attn(inputNorm(dx)) + } + if e = firstErr(e, d.recordFFNSublayer(r, b, 1)); e != nil { // MoE (raw softmax) or dense SwiGLU + r.Free() + return nil, e + } + continue + } + if d.mla { // DeepSeek-V2 Multi-head Latent Attention (rectangular, low-rank latent KV) + if e := firstErr(d.recordMLAAttention(r, b, pos, 1), d.recordFFNSublayer(r, b, 1)); e != nil { + r.Free() + return nil, e + } + continue + } + // attention projections: fused single QKV matmul when available (§T613 — the q/k + // bands rotate IN PLACE inside the combined row, k/v append to the cache straight + // from their bands, attention reads q at band offset 0), else the unfused three. + // Recording order is execution order: norm FIRST, then the projection branch. + e := d.recordAttnNorm(r, b, 1) + e = firstErr(e, recordBarrier(r)) + qBuf := d.q.b + if e == nil && b.wqkv != nil { + qBuf = d.qkv.b + e = firstErr( + d.recordQKVProj(r, b, 1), + d.recordQKNorm(r, b, 1, D+2*kvDim), // per-head QK-norm (Qwen3); no-op otherwise + ) + e = firstErr(e, recordBarrier(r)) + if fr, ok := r.(ropePairF16KVAppendRecorder); e == nil && ok && d.f16KV && + !d.noRope && d.rotaryDim == dk && dk == 64 { + e = fr.RoPEPairF16KVAppend( + d.qkv.b, d.dinv.b, b.kC, b.vC, D+2*kvDim, + H, 0, KVH, D, dk, d.half, D+kvDim, kvDim, pos, pos*kvDim, d.posDiv, + ) + } else if e == nil { + e = firstErr( + d.ropeQK(r, d.qkv.b, d.dinv.b, 1, D+2*kvDim, pos), // q+k bands, ONE dispatch (§T613); full ∨ partial rotary + d.recordKVPairBlit(r, d.qkv.b, D, d.qkv.b, D+kvDim, b.kC, b.vC, pos*kvDim, kvDim), + ) + } + } else if e == nil { + e = firstErr( + b.wq.record(r, d.xn.b, d.q.b, 1), + b.wk.record(r, d.xn.b, d.k.b, 1), + b.wv.record(r, d.xn.b, d.v_.b, 1), + ) + e = firstErr(e, recordBarrier(r)) + if fr, ok := r.(ropeF16KVAppendRecorder); e == nil && ok && d.f16KV && + !d.noRope && d.rotaryDim == dk && dk == 64 { + e = fr.RoPEF16KVAppend( + d.q.b, d.k.b, d.v_.b, d.dinv.b, b.kC, b.vC, + H, KVH, dk, d.half, pos, pos*kvDim, d.posDiv, + ) + } else if e == nil { + e = firstErr( + r.RoPE(d.q.b, d.dinv.b, d.q.b, 1, D, H, dk, d.half, pos, d.posDiv), + r.RoPE(d.k.b, d.dinv.b, d.k.b, 1, kvDim, KVH, dk, d.half, pos, d.posDiv), + d.recordKVPairBlit(r, d.k.b, 0, d.v_.b, 0, b.kC, b.vC, pos*kvDim, kvDim), + ) + } + } + e = firstErr(e, recordBarrier(r)) + e = firstErr(e, d.recordMHA(r, qBuf, b.kC, b.vC, d.attn.b, 1, pos+1)) + e = firstErr(e, recordBarrier(r)) + e = firstErr(e, d.recordOProj(r, b, 1)) // dx += attn·Wo (+ optional o-bias) + e = firstErr(e, recordBarrier(r)) + e = firstErr(e, d.recordFFNSublayer(r, b, 1)) + e = firstErr(e, recordBarrier(r)) + if e != nil { + r.Free() + return nil, e + } + } + e := d.norm(r, d.dx.b, d.gFinal.b, d.bFinalBeta(), d.xn.b, 1) + e = firstErr(e, recordBarrier(r)) + e = firstErr(e, d.recordLogits(r, d.logits.b, 1)) + if e != nil { + r.Free() + return nil, e + } + return r, nil +} + +// dropPending frees a pre-encoded next-step recorder that can no longer be used (the +// position moved unexpectedly, a prefill rewrote the cache, or the decoder is released). +func (d *Decoder) dropPending() { + if d.pending != nil { + d.pending.Free() + d.pending = nil + } + d.pendingPos = -1 +} + +// Step advances the decoder by one token at absolute position pos (== cache length before the call), +// recording the whole layer stack + vocab head into one command buffer, and returns the [vocab] +// logits. pos must be < the model's Ctx. Sequential calls run the §T614 encode-overlap: while the +// GPU executes this step, the next position's command buffer is pre-encoded on the host. +func (d *Decoder) stepInto(token, pos int) error { + if pos < 0 || pos >= d.maxLen { + return fmt.Errorf("llamagpu(%s): pos %d out of [0,%d)", d.ops.name, pos, d.maxLen) + } + if token < 0 || token >= d.v { + return fmt.Errorf("llamagpu(%s): token %d out of vocab %d", d.ops.name, token, d.v) + } + // Mamba/Mamba-2/RWKV carry per-block conv/SSM state across Step calls with no KV cache; a fresh + // sequence (pos==0) must start from zeroed state. + if pos == 0 { + if d.mamba { + if err := d.resetMambaState(); err != nil { + return err + } + } + if d.mamba2 { + if err := d.resetMamba2State(); err != nil { + return err + } + } + if d.rwkv { + if err := d.resetRWKVState(); err != nil { + return err + } + } + if d.jamba { // reset the Mamba-layer conv/SSM state (attention layers' KV cache overwrites from row 0) + if err := d.resetMambaState(); err != nil { + return err + } + } + } + // the token's embedding must be resident BEFORE commit — the recorded chain reads d.dx. + if err := d.dx.b.UploadF32(d.gatherEmbed(token)); err != nil { + return err + } + var r recorder + if d.pending != nil && d.pendingPos == pos { + r, d.pending, d.pendingPos = d.pending, nil, -1 + } else { + d.dropPending() + var err error + if r, err = d.encodeStep(pos); err != nil { + return err + } + } + if err := r.Commit(); err != nil { + r.Free() + return err + } + // encode-overlap (§T614): pre-encode pos+1 while the GPU executes pos. Failure is not fatal — the + // next Step simply encodes fresh. Disabled for Mamba: its blocks mutate resident conv/SSM state, so + // a pre-encoded pos+1 must not be recorded until pos's state writes are committed. + if d.ops.asyncEncode && !d.mamba && !d.mamba2 && !d.rwkv && pos+1 < d.maxLen { + if nr, err := d.encodeStep(pos + 1); err == nil { + d.pending, d.pendingPos = nr, pos+1 + } + } + if err := r.Wait(); err != nil { + r.Free() + return err + } + r.Free() + return nil +} + +// Step records and executes one decode step for token at pos, returning the host logits [d.v]. It wraps +// stepInto (which leaves the result in the device logits buffer) with the whole-vocab host download; the +// on-device sampling fast-path (Generate) instead reads the device logits directly via a top-k. +func (d *Decoder) Step(token, pos int) ([]float32, error) { + if err := d.stepInto(token, pos); err != nil { + return nil, err + } + out := make([]float32, d.v) + if err := d.logits.b.DownloadF32(out); err != nil { + return nil, err + } + return out, nil +} + +// StepN advances the decoder by k tokens at absolute positions pos..pos+k-1 in ONE recorded step +// (§T418): the whole layer stack runs over [k,·] rows with causal attention against the growing +// cache (row i attends up to pos+i), and the new k KV rows are appended. Returns the [k,vocab] +// logits (row i = logits after tokens[..i]). This is the prompt-PREFILL fast path — one dispatch +// round-trip for the whole prompt instead of one per token — and the target-verification step +// speculative decoding needs. pos+k must be ≤ the model's Ctx. +func (d *Decoder) StepN(tokens []int, pos int) ([]float32, error) { return d.stepN(tokens, pos, false) } + +// StepNLast is StepN but projects and downloads ONLY the final prompt row's logits (the LM head runs +// over 1 row instead of k, and the host download shrinks from k·vocab to vocab). The transformer body +// still runs all k rows to populate the KV cache, so the returned [vocab] row is bit-identical to the +// last row of StepN. Use it for prefill whose only consumer is the post-prompt token (generation) or +// which discards logits entirely (speculative/prompt-lookup KV-seeding) — the [k-1] earlier rows' LM +// head + PCIe download are pure waste there (a 512-token prompt over a 128k vocab downloads 256 MB just +// to discard all but the last row). StepN (full-k) is kept for verify callers that need every row. +func (d *Decoder) StepNLast(tokens []int, pos int) ([]float32, error) { + return d.stepN(tokens, pos, true) +} + +// stepN is StepN with an optional last-row-only mode. When lastOnly, only the FINAL prompt row's +// logits are projected by the LM head and downloaded — the transformer body still runs all k rows +// (the KV cache needs every position), but the [k-1] earlier rows' logits are never read by a caller +// that only samples the post-prompt token (Generate), so their lm_head GEMM and vocab-sized host +// download are pure waste — large for high-vocab models (a 512-token prompt over a 128k vocab downloads +// 256 MB just to discard all but the last row). It then returns a single [vocab] row. Full-k mode is +// unchanged and keeps every row for speculative/verify callers (Medusa, prompt-lookup). +func (d *Decoder) stepN(tokens []int, pos int, lastOnly bool) ([]float32, error) { + k := len(tokens) + if k == 0 { + return nil, fmt.Errorf("llamagpu(%s): StepN needs ≥1 token", d.ops.name) + } + if pos < 0 || pos+k > d.maxLen { + return nil, fmt.Errorf("llamagpu(%s): StepN [%d,%d) out of [0,%d)", d.ops.name, pos, pos+k, d.maxLen) + } + if d.mamba || d.mamba2 || d.rwkv || d.jamba { + // Mamba/Mamba-2/RWKV/Jamba's recurrence is inherently sequential (each token's state update reads the + // previous token's), so there is no batched prefill — process the prompt token by token through + // Step, which carries the recurrent state and resets it at pos==0. Returns [k, vocab]. + if lastOnly { // only the final token's logits are wanted; still step all tokens for the state + var last []float32 + for i, tok := range tokens { + row, err := d.Step(tok, pos+i) + if err != nil { + return nil, err + } + last = row + } + return append([]float32(nil), last...), nil + } + out := make([]float32, k*d.v) + for i, tok := range tokens { + row, err := d.Step(tok, pos+i) + if err != nil { + return nil, err + } + copy(out[i*d.v:(i+1)*d.v], row) + } + return out, nil + } + logits := d.logits.b + if !lastOnly { + var err error + logits, err = logitsForRows(d.ops, d.logits.b, d.logitsRows, &d.fullLogits, k, d.v) + if err != nil { + return nil, err + } + } + // a batched step rewrites the cache and scratch — any single-step encode pre-staged for + // the old position is now invalid (§T614). + d.dropPending() + host := make([]float32, k*d.d) + for i, tok := range tokens { + if tok < 0 || tok >= d.v { + return nil, fmt.Errorf("llamagpu(%s): token %d out of vocab %d", d.ops.name, tok, d.v) + } + copy(host[i*d.d:(i+1)*d.d], d.gatherEmbed(tok)) + } + if err := d.dx.b.UploadF32(host); err != nil { + return nil, err + } + r, err := d.ops.newRecorder() + if err != nil { + return nil, err + } + D, H, KVH, dk, kvDim := d.qDim, d.h, d.kvH, d.dk, d.kvDim // D = attention Q/O width (h·dk); = d.d unless head_dim decoupled + stride := D + 2*kvDim + for _, b := range d.blocks { + if d.mla { // DeepSeek-V2 MLA prefill: k new tokens at position pos + if e := firstErr(d.recordMLAAttention(r, b, pos, k), d.recordFFNSublayer(r, b, k)); e != nil { + r.Free() + return nil, e + } + continue + } + // fused QKV (§T613): one [k,stride] matmul; q/k bands rotate in place (RoPEAt's + // width parameter acts as the row stride), then Copy2D extracts q contiguously + // and deposits the k/v bands directly as cache rows. Recording order = execution + // order: norm first, then the projection branch. + e := d.recordAttnNorm(r, b, k) + attnK, attnV, attnF16 := b.kC, b.vC, d.f16KV + prefillQKV := b.wqkv + mixedQKV := prefillQKV == nil && b.wqkvPrefill != nil && k >= 24 + if mixedQKV { + prefillQKV = b.wqkvPrefill + } + if e == nil && prefillQKV != nil { + e = firstErr( + d.recordQKVProjWith(r, prefillQKV, b.qkvBias, k), + d.recordQKNorm(r, b, k, stride), // per-head QK-norm (Qwen3); no-op otherwise + ) + split := false + if sr, ok := r.(ropePairSplitRecorder); ok && mixedQKV && !d.noRope && d.rotaryDim == d.dk { + // The mixed group wins at the projection but lost its full-model gate when the + // canonical fused layout was unpacked by three standalone Copy2D dispatches. Rotate + // q/k and scatter q/k/v together; cache conversion then reads contiguous k/v. + e = firstErr(e, sr.RoPEPairSplit( + d.qkv.b, d.dinv.b, d.q.b, d.k.b, d.v_.b, + k, stride, H, 0, KVH, D, dk, d.half, D+kvDim, kvDim, pos, d.posDiv, + )) + split = true + if d.f16KV && pos == 0 { + attnK, attnV, attnF16 = d.k.b, d.v_.b, false + } + e = firstErr(e, d.recordKVPairBlit( + r, d.k.b, 0, d.v_.b, 0, b.kC, b.vC, pos*kvDim, k*kvDim, + )) + } + if !split { + e = firstErr( + e, + d.ropeQK(r, d.qkv.b, d.dinv.b, k, stride, pos), // q+k bands, ONE dispatch (§T613); full ∨ partial rotary + r.Copy2D(d.qkv.b, 0, stride, d.q.b, 0, D, k, D), + ) + // At initial prefill, retain contiguous f32 K/V scratch for the established FlashMM + // attention path while independently converting the same rows into the f16 decode cache. + if d.f16KV && pos == 0 { + e = firstErr(e, + r.Copy2D(d.qkv.b, D, stride, d.k.b, 0, kvDim, k, kvDim), + r.Copy2D(d.qkv.b, D+kvDim, stride, d.v_.b, 0, kvDim, k, kvDim), + ) + attnK, attnV, attnF16 = d.k.b, d.v_.b, false + } + e = firstErr(e, d.recordKVPairCopy2D( + r, d.qkv.b, D, stride, d.qkv.b, D+kvDim, stride, + b.kC, b.vC, pos*kvDim, kvDim, k, kvDim, + )) + } + } else if e == nil { + e = firstErr( + b.wq.record(r, d.xn.b, d.q.b, k), + b.wk.record(r, d.xn.b, d.k.b, k), + b.wv.record(r, d.xn.b, d.v_.b, k), + r.RoPE(d.q.b, d.dinv.b, d.q.b, k, D, H, dk, d.half, pos, d.posDiv), + r.RoPE(d.k.b, d.dinv.b, d.k.b, k, kvDim, KVH, dk, d.half, pos, d.posDiv), + d.recordKVPairBlit(r, d.k.b, 0, d.v_.b, 0, b.kC, b.vC, pos*kvDim, k*kvDim), + ) + if d.f16KV && pos == 0 { + attnK, attnV, attnF16 = d.k.b, d.v_.b, false + } + } + e = firstErr( + e, + // sq=k vs sk=pos+k: the kernel's causal offset (sk-sq = pos) makes row i attend + // through absolute position pos+i — exactly the prefill/verify semantics. + d.recordMHAStorage(r, d.q.b, attnK, attnV, d.attn.b, k, pos+k, attnF16), + d.recordOProj(r, b, k), // dx += attn·Wo (+ optional o-bias) + d.recordFFNSublayer(r, b, k), + ) + if e != nil { + r.Free() + return nil, e + } + } + rows := k + if lastOnly { + // Only the final row's logits are wanted: move its residual to row 0 (r.Blit copies d.d floats + // from d.dx[(k-1)·d] to d.dx[0]; a no-op when k==1) so the final norm + LM-head project ONE row + // instead of k, and the host download shrinks from k·vocab to vocab. Bit-identical: the discarded + // rows' logits were never read, and row 0's value is exactly what row k-1 would have produced. + if e := r.Blit(d.dx.b, (k-1)*d.d, d.dx.b, 0, d.d); e != nil { + r.Free() + return nil, e + } + rows = 1 + } + if e := firstErr( + d.norm(r, d.dx.b, d.gFinal.b, d.bFinalBeta(), d.xn.b, rows), + d.recordLogits(r, logits, rows), + r.Finish(), + ); e != nil { + r.Free() + return nil, e + } + r.Free() + out := make([]float32, rows*d.v) + if err := logits.DownloadF32(out); err != nil { + return nil, err + } + return out, nil +} + +// Generate runs the batched decode as a text-generation loop: it prefills the prompt, then samples +// up to maxNew tokens (bounded by the model's Ctx), feeding each back. Returns prompt+generated +// token ids. With a greedy sampler it produces the same ids as nlp.Llama.Generate. +func (d *Decoder) Generate(prompt []int, maxNew int, s nlp.TokenSampler) ([]int, error) { + if len(prompt) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Generate needs a non-empty prompt", d.ops.name) + } + out := append([]int(nil), prompt...) + // prefill the whole prompt in ONE recorded step (§T418) — one dispatch round-trip instead of one + // per prompt token. lastOnly: generation samples only the post-prompt row, so the LM head projects + // and downloads that single row (not all len(prompt) rows) — a large saving for high-vocab models. + all, err := d.stepN(prompt, 0, true) + if err != nil { + return nil, err + } + pos := len(prompt) + logits := all // last-row-only prefill returns exactly the post-prompt row's [vocab] logits + buf := make([]float64, d.v) + // Device-resident top-k sampling fast-path: for a penalty-free TopK sampler AND a logits buffer that + // can TopK (CUDA GPU reduction or Metal coherent-UMA selection), skip the whole-vocab host download. + // Read the K highest logits at the resident boundary and draw over them, bit-identical to the + // full-vocab sampler (see fastTopKSampler/sampleTopKCandidates). Any other sampler/backend takes the + // flexible host path. + sp, fastK := fastTopKSampler(s, d.v) + spP, fastC := fastTopPSampler(s, d.v) // pure-top-p device path (mutually exclusive with sp: needs top-k off) + dt, hasDev := d.logits.b.(deviceTopKer) + dtP, hasDevP := d.logits.b.(deviceTopPer) + fast := (sp != nil && hasDev) || (spP != nil && hasDevP) + for range maxNew { + if pos >= d.maxLen { + break + } + var next int + if fast && logits == nil { // decode step, logits resident on device (first d.v of the buffer) + if sp != nil { // top-k (± post-shrink top-p/min-p): draw over the K-candidate superset + gi, gv, e := dt.TopKN(d.v, fastK) + if e != nil { + return nil, e + } + next = sampleTopKCandidates(sp, gi, gv) + } else { // pure top-p: resolve the nucleus over the top-C candidates using full-vocab stats + maxL, zexp, e2 := dtP.SoftmaxStatsN(d.v, spP.Temperature) + if e2 != nil { + return nil, e2 + } + resolved := false + // Cheap overflow guard (stats only): top-C mass ≤ C/Zexp; below TopP the nucleus can't fit, + // so skip the O(n·C) device TopK and fall straight to the host path (see the graph decoder). + if float64(fastC)/zexp >= spP.TopP { + gi, gv, e := dtP.TopKN(d.v, fastC) + if e != nil { + return nil, e + } + cl := make([]float64, len(gv)) + for i, v := range gv { + cl[i] = float64(v) + } + if tok, ok := spP.SampleTopPFromCandidates(cl, gi, maxL, zexp); ok { + next = tok + resolved = true + } + } + if !resolved { // guard predicted overflow, or TopK confirmed it → full-vocab host sample + l, e3 := dtP.ToHost() + if e3 != nil { + return nil, e3 + } + lf := l.Storage().F32() + for i := 0; i < d.v; i++ { + buf[i] = float64(lf[i]) + } + next = s.SampleWithHistory(buf, out) + } + } + } else { // first token (prefill host logits) or the full-vocab fallback + for i, x := range logits { + buf[i] = float64(x) + } + next = s.SampleWithHistory(buf, out) + } + out = append(out, next) + if fast { + if e := d.stepInto(next, pos); e != nil { // leaves logits on device — no host download + return nil, e + } + logits = nil + } else { + l, e := d.Step(next, pos) + if e != nil { + return nil, e + } + logits = l + } + pos++ + } + return out, nil +} + +// Vocab returns the model's vocabulary size. +func (d *Decoder) Vocab() int { return d.v } + +// Ctx returns the model's maximum context length (the KV-cache capacity). +func (d *Decoder) Ctx() int { return d.maxLen } + +// Release frees all device buffers and resident quantized weights. +func (d *Decoder) Release() { + d.dropPending() + d.fullLogits.release() + for _, b := range d.all { + if b != nil { + b.Release() + } + } + d.all = nil + for _, w := range d.qweights { + if w != nil { + _ = w.Close() + } + } + d.qweights = nil +} + +func boolToInt(b bool) int { + if b { + return 1 + } + return 0 +} + +func firstErr(errs ...error) error { + for _, e := range errs { + if e != nil { + return e + } + } + return nil +} + +// rowBlock is the number of source rows converted at a time by forEachRowBlockF32; colTile is the +// column tile flat2DT transposes within such a block. Together they make the transpose a classic +// cache-blocked 2D tile walk: both the contiguous reads and the strided writes stay inside a +// cache-resident working set instead of touching a fresh line (and TLB entry) per element. The +// pair was picked by measurement on a [32000, 2048] lm_head — see the §V22 benchmarks in +// flatten_test.go, where tiling is worth ~1.8× on top of the bulk-slice conversion alone. +const ( + rowBlock = 256 + colTile = 64 +) + +// forEachRowBlockF32 walks t (rank-2, any dtype, any strides) in blocks of at most rowBlock rows, +// calling fn with the block's first row index and its values converted to float32 in contiguous +// row-major order (len(blk) == (j1-j0)·cols). +// +// This is the bulk-slice idiom the base-perf sweep established: one dtype decision per BLOCK via +// tensor.Cast's typed fast paths, instead of one AtF64 interface dispatch per ELEMENT. It is +// numerically identical to float32(t.AtF64(j, i)) for every dtype — Cast's F64→F32, F16→F32 and +// BF16→F32 paths reproduce the widen-then-narrow composition exactly — so callers stay +// bit-identical to the per-element loops they replace. +func forEachRowBlockF32(t *tensor.Tensor, fn func(j0 int, blk []float32)) { + rows, cols := t.Shape()[0], t.Shape()[1] + //perfscan:ignore PS2001 Cast alloc per row-BLOCK (256), model-build bulk-slice upload, amortized + for j0 := 0; j0 < rows; j0 += rowBlock { + j1 := min(j0+rowBlock, rows) + blk, err := t.Slice(0, j0, j1) + if err != nil { // unreachable for in-range slices; degrade to the per-element read + scratch := make([]float32, (j1-j0)*cols) + for j := j0; j < j1; j++ { + for i := range cols { + scratch[(j-j0)*cols+i] = float32(t.AtF64(j, i)) + } + } + fn(j0, scratch) + continue + } + if blk.Dtype() == tensor.F32 && blk.IsContiguous() { + // Already the target layout and dtype: hand out the storage window itself, no + // conversion and no allocation. This is the f32-weight upload path. + off := blk.Offset() + fn(j0, blk.Storage().F32()[off:off+(j1-j0)*cols]) + continue + } + fn(j0, blk.Cast(tensor.F32).Storage().F32()[:(j1-j0)*cols]) + } +} + +// flat2DT flattens the TRANSPOSE of a rank-2 tensor into a row-major float32 slice: +// out[i·rows+j] == float32(t.AtF64(j, i)) for t of shape [rows, cols], i.e. the [cols, rows] +// layout the device matmuls expect from a torch-style [out, in] weight or a tied lm_head. +// +// Used for weight upload at model-build time, where the tensor can be huge — Gemma's tied lm_head +// is [256000, 2048] = 524M elements — so this runs as a blocked bulk-slice transpose rather than +// per-element reads. +func flat2DT(t *tensor.Tensor) []float32 { + o := make([]float32, t.Shape()[0]*t.Shape()[1]) + flat2DTInto(o, t) + return o +} + +// flat2DTInto is flat2DT writing into a caller-provided destination of exactly rows·cols floats. +// It lets a caller that already owns the target buffer — e.g. the tied-lm_head constructors, which +// need the result inside a [dim, vocab] tensor — skip an intermediate slice, which at Gemma's +// 256k vocab is a 2 GB allocation and a 2 GB copy saved. +func flat2DTInto(o []float32, t *tensor.Tensor) { + rows, cols := t.Shape()[0], t.Shape()[1] + forEachRowBlockF32(t, func(j0 int, blk []float32) { + j1 := j0 + len(blk)/cols + for i0 := 0; i0 < cols; i0 += colTile { + i1 := min(i0+colTile, cols) + for j := j0; j < j1; j++ { + row := blk[(j-j0)*cols+i0 : (j-j0)*cols+i1] + for i, v := range row { + o[(i0+i)*rows+j] = v + } + } + } + }) +} + +// flat2D flattens a rank-2 tensor into a row-major float32 slice (no transpose): +// out[i·cols+j] == float32(t.AtF64(i, j)). +func flat2D(t *tensor.Tensor) []float32 { + r, c := t.Shape()[0], t.Shape()[1] + o := make([]float32, r*c) + forEachRowBlockF32(t, func(j0 int, blk []float32) { copy(o[j0*c:], blk) }) + return o +} + +// flat1D flattens a rank-1 tensor (a norm γ, a bias) into a float32 slice. +func flat1D(t *tensor.Tensor) []float32 { + n := t.Shape()[0] + o := make([]float32, n) + copy(o, t.Cast(tensor.F32).Storage().F32()[:n]) + return o +} + +// embedRow reads one row of an embedding table as float32. It is on the per-token decode path, so +// it takes the bulk-slice route too: one row view, one typed conversion, no per-element dispatch. +// +//perfscan:ignore PS3033 one small row alloc/token, resource-only, negligible vs forward +func embedRow(table *tensor.Tensor, row, cols int) []float32 { + o := make([]float32, cols) + r, err := table.Slice(0, row, row+1) + if err != nil { + for j := range cols { + o[j] = float32(table.AtF64(row, j)) + } + return o + } + copy(o, r.Cast(tensor.F32).Storage().F32()[:cols]) + return o +} + +// gatherEmbed reads a token's embedding row and applies the Granite EmbeddingMult (embMult == 1 +// for every non-Granite model, so this is embedRow verbatim then). +func (d *Decoder) gatherEmbed(token int) []float32 { + e := embedRow(d.table, token, d.d) + if d.embMult != 1 { + for i := range e { + e[i] *= d.embMult + } + } + return e +} diff --git a/checks/testdata/ps6136-owner/after/gpt.go.txt b/checks/testdata/ps6136-owner/after/gpt.go.txt new file mode 100644 index 00000000..3f70fb8b --- /dev/null +++ b/checks/testdata/ps6136-owner/after/gpt.go.txt @@ -0,0 +1,404 @@ +package llamagpu + +import ( + "fmt" + "math" + + "github.com/jxsl13/goai/nlp" + "github.com/jxsl13/goai/tensor" +) + +// gptBlock is one pre-LN GPT transformer block on the device: LayerNorm gains/biases, the four +// attention projections, the biased GELU MLP, and the block's KV cache. +type gptBlock struct { + g1, b1, g2, b2 buffer // LN1/LN2 gamma+beta [d] + wq, wk, wv, wo linear // attention projections [d,d], no bias + wqkv f32Linear // fused attention projection [d,3d] + w1, w2 linear // FFN up [d,4d] / down [4d,d] + fb1, fb2 buffer // FFN biases [4d] / [d] + kC, vC buffer // KV caches [maxLen*d] + ffn int // FFN width, fixed when the block is uploaded +} + +// GPTDecoder is the batched-decode engine for nlp.GPT models (§T422) — the GPT-2-style sibling of +// [Decoder]: LayerNorm instead of RMSNorm, learned positional embeddings instead of RoPE, a biased +// GELU MLP instead of SwiGLU. Each Step records the whole layer stack into one command buffer. +// Not safe for concurrent use; Release when done. +type GPTDecoder struct { + ops backendOps + d, h, dk, v, maxLen int + eps, scale float32 + + blocks []gptBlock + lnfG, lnfB *bufSlot + head linear + dx, xn, xn2 *bufSlot + q, k, v_, qkv *bufSlot + attn, hid *bufSlot + logits *bufSlot + logitsRows int + fullLogits growBuffer + tokEmb, posEmb *tensor.Tensor // host gather sources + all []buffer +} + +type f32QKVBandsRecorder interface { + F32QKVBands(x, w, q, k, v buffer, rows, dim int) error +} + +type biasGELURecorder interface { + BiasGELU(x, b, o buffer, rows, n int) error +} + +func (d *GPTDecoder) recordAttentionResidual(r recorder, b gptBlock, rows int) error { + return b.wo.recordAdd(r, d.attn.b, nil, d.dx.b, rows, d.d) +} + +func (d *GPTDecoder) recordFFNResidual(r recorder, b gptBlock, rows int) error { + return firstErr( + b.w2.recordAdd(r, d.hid.b, nil, d.dx.b, rows, d.d), + r.AddBias(d.dx.b, b.fb2, d.dx.b, rows, d.d), + ) +} + +func (d *GPTDecoder) recordBiasGELU(r recorder, b gptBlock, rows int) error { + if d.ops.fusedBiasGELU { + if fused, ok := r.(biasGELURecorder); ok { + return fused.BiasGELU(d.hid.b, b.fb1, d.hid.b, rows, b.ffn) + } + } + return firstErr( + r.AddBias(d.hid.b, b.fb1, d.hid.b, rows, b.ffn), + r.Unary(d.hid.b, d.hid.b, unaryGELU), + ) +} + +// newGPTDecoder uploads m's weights via ops and prepares per-block KV caches up to Ctx. +func newGPTDecoder(m *nlp.GPT, ops backendOps) (*GPTDecoder, error) { + cfg := m.Config + d := &GPTDecoder{ + ops: ops, + d: cfg.Dim, h: cfg.Heads, v: cfg.Vocab, maxLen: cfg.Ctx, + eps: float32(cfg.Eps), tokEmb: m.TokEmb, posEmb: m.PosEmb, + } + if d.h <= 0 || d.d%d.h != 0 { + return nil, fmt.Errorf("llamagpu(%s): gpt dim %d not divisible by heads %d", ops.name, d.d, d.h) + } + d.dk = d.d / d.h + d.scale = float32(1.0 / math.Sqrt(float64(d.dk))) + + var err error + mk := func(data []float32) *bufSlot { + if err != nil { + return &bufSlot{} + } + b, e := ops.newBuffer(data) + if e != nil { + err = e + return &bufSlot{} + } + d.all = append(d.all, b) + return &bufSlot{b: b} + } + lin := func(w *tensor.Tensor) linear { + in, out := w.Shape()[0], w.Shape()[1] + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } + fusedQKV := func(wq, wk, wv *tensor.Tensor) f32Linear { + in := wq.Shape()[0] + nq, nk, nv := wq.Shape()[1], wk.Shape()[1], wv.Shape()[1] + fq, fk, fv := flat2D(wq), flat2D(wk), flat2D(wv) + stride := nq + nk + nv + weight := make([]float32, in*stride) + for row := range in { + dst := weight[row*stride : (row+1)*stride] + copy(dst[:nq], fq[row*nq:(row+1)*nq]) + copy(dst[nq:nq+nk], fk[row*nk:(row+1)*nk]) + copy(dst[nq+nk:], fv[row*nv:(row+1)*nv]) + } + return f32Linear{w: mk(weight).b, k: in, n: stride} + } + for _, b := range m.Blocks { + gb := gptBlock{ + g1: mk(flat1D(b.LN1.Gamma)).b, b1: mk(flat1D(b.LN1.Beta)).b, + g2: mk(flat1D(b.LN2.Gamma)).b, b2: mk(flat1D(b.LN2.Beta)).b, + wo: lin(b.Attn.Wo), w1: lin(b.W1), w2: lin(b.W2), + fb1: mk(flat1D(b.B1)).b, fb2: mk(flat1D(b.B2)).b, + kC: mk(make([]float32, d.maxLen*d.d)).b, vC: mk(make([]float32, d.maxLen*d.d)).b, + ffn: b.W1.Shape()[1], + } + if ops.fusedF32QKV { + gb.wqkv = fusedQKV(b.Attn.Wq, b.Attn.Wk, b.Attn.Wv) + } else { + gb.wq, gb.wk, gb.wv = lin(b.Attn.Wq), lin(b.Attn.Wk), lin(b.Attn.Wv) + } + d.blocks = append(d.blocks, gb) + } + d.lnfG = mk(flat1D(m.LNf.Gamma)) + d.lnfB = mk(flat1D(m.LNf.Beta)) + d.head = lin(m.Head) + ffn := 4 * d.d + if len(m.Blocks) > 0 { + ffn = m.Blocks[0].W1.Shape()[1] + } + c := d.maxLen // scratch sized for multi-token StepN (§T423) + d.dx = mk(make([]float32, c*d.d)) + d.xn = mk(make([]float32, c*d.d)) + d.xn2 = mk(make([]float32, c*d.d)) + d.q = mk(make([]float32, c*d.d)) + d.k = mk(make([]float32, c*d.d)) + d.v_ = mk(make([]float32, c*d.d)) + if ops.fusedF32QKV { + // Large prefills address bands of the single fused weight directly. Only rows below the + // strided-band crossover need a grouped-output scratch, bounding it to 63 rows. + d.qkv = mk(make([]float32, min(c, 63)*3*d.d)) + } + d.attn = mk(make([]float32, c*d.d)) + d.hid = mk(make([]float32, c*ffn)) + d.logitsRows = 1 + if ops.eagerFullLogits { + d.logitsRows = c + } + d.logits = mk(make([]float32, d.logitsRows*d.v)) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// Step advances the decoder by one token at absolute position pos (== cache length before the +// call), recording the whole layer stack + head into one command buffer, and returns the [vocab] +// logits. The token embedding and learned positional embedding are summed host-side. +func (d *GPTDecoder) Step(token, pos int) ([]float32, error) { + if pos < 0 || pos >= d.maxLen { + return nil, fmt.Errorf("llamagpu(%s): gpt pos %d out of [0,%d)", d.ops.name, pos, d.maxLen) + } + if token < 0 || token >= d.v { + return nil, fmt.Errorf("llamagpu(%s): gpt token %d out of vocab %d", d.ops.name, token, d.v) + } + x := embedRow(d.tokEmb, token, d.d) + pe := embedRow(d.posEmb, pos, d.d) + for i := range x { + x[i] += pe[i] + } + if err := d.dx.b.UploadF32(x); err != nil { + return nil, err + } + r, err := d.ops.newRecorder() + if err != nil { + return nil, err + } + D, H, dk := d.d, d.h, d.dk + for _, b := range d.blocks { + e := r.LayerNorm(d.dx.b, b.g1, b.b1, d.xn.b, 1, D, d.eps) + qBuf := d.q.b + if d.ops.fusedF32QKV { + qBuf = d.qkv.b + e = firstErr(e, + b.wqkv.record(r, d.xn.b, d.qkv.b, 1), + r.Blit(d.qkv.b, D, b.kC, pos*D, D), + r.Blit(d.qkv.b, 2*D, b.vC, pos*D, D), + ) + } else { + e = firstErr(e, + b.wq.record(r, d.xn.b, d.q.b, 1), + b.wk.record(r, d.xn.b, d.k.b, 1), + b.wv.record(r, d.xn.b, d.v_.b, 1), + r.Blit(d.k.b, 0, b.kC, pos*D, D), + r.Blit(d.v_.b, 0, b.vC, pos*D, D), + ) + } + e = firstErr(e, + r.MHA(qBuf, b.kC, b.vC, d.attn.b, 1, pos+1, D, H, H, dk, 1, 0, d.scale), + d.recordAttentionResidual(r, b, 1), + r.LayerNorm(d.dx.b, b.g2, b.b2, d.xn2.b, 1, D, d.eps), + b.w1.record(r, d.xn2.b, d.hid.b, 1), + d.recordBiasGELU(r, b, 1), + d.recordFFNResidual(r, b, 1), + ) + if e != nil { + r.Free() + return nil, e + } + } + if e := firstErr( + r.LayerNorm(d.dx.b, d.lnfG.b, d.lnfB.b, d.xn.b, 1, D, d.eps), + d.head.record(r, d.xn.b, d.logits.b, 1), + r.Finish(), + ); e != nil { + r.Free() + return nil, e + } + r.Free() + out := make([]float32, d.v) + if err := d.logits.b.DownloadF32(out); err != nil { + return nil, err + } + return out, nil +} + +// StepN advances the decoder by k tokens at absolute positions pos..pos+k-1 in ONE recorded step +// (§T423, the GPT sibling of Decoder.StepN): the layer stack runs over [k,·] rows with causal +// attention against the growing cache, biases broadcast via the AddBias record op, and the k KV +// rows are appended. Returns the [k,vocab] logits. +func (d *GPTDecoder) StepN(tokens []int, pos int) ([]float32, error) { + return d.gptStepN(tokens, pos, false) +} + +// StepNLast is StepN projecting only the FINAL prompt row's logits (LM head over 1 row, download of +// vocab not k·vocab); the transformer body still runs all k rows so the KV cache and the returned +// [vocab] row are bit-identical to StepN's last row. For generation/KV-seeding prefills that read at +// most the last row — mirrors Decoder.StepNLast. +func (d *GPTDecoder) StepNLast(tokens []int, pos int) ([]float32, error) { + return d.gptStepN(tokens, pos, true) +} + +func (d *GPTDecoder) gptStepN(tokens []int, pos int, lastOnly bool) ([]float32, error) { + k := len(tokens) + if k == 0 { + return nil, fmt.Errorf("llamagpu(%s): gpt StepN needs ≥1 token", d.ops.name) + } + if pos < 0 || pos+k > d.maxLen { + return nil, fmt.Errorf("llamagpu(%s): gpt StepN [%d,%d) out of [0,%d)", d.ops.name, pos, pos+k, d.maxLen) + } + logits := d.logits.b + if !lastOnly { + var err error + logits, err = logitsForRows(d.ops, d.logits.b, d.logitsRows, &d.fullLogits, k, d.v) + if err != nil { + return nil, err + } + } + host := make([]float32, k*d.d) + for i, tok := range tokens { + if tok < 0 || tok >= d.v { + return nil, fmt.Errorf("llamagpu(%s): gpt token %d out of vocab %d", d.ops.name, tok, d.v) + } + row := host[i*d.d : (i+1)*d.d] + copy(row, embedRow(d.tokEmb, tok, d.d)) + pe := embedRow(d.posEmb, pos+i, d.d) + for j := range row { + row[j] += pe[j] + } + } + if err := d.dx.b.UploadF32(host); err != nil { + return nil, err + } + r, err := d.ops.newRecorder() + if err != nil { + return nil, err + } + D, H, dk := d.d, d.h, d.dk + for _, b := range d.blocks { + e := r.LayerNorm(d.dx.b, b.g1, b.b1, d.xn.b, k, D, d.eps) + if d.ops.fusedF32QKV { + if bands, ok := r.(f32QKVBandsRecorder); ok && k >= 64 { + e = firstErr(e, bands.F32QKVBands(d.xn.b, b.wqkv.w, d.q.b, d.k.b, d.v_.b, k, D)) + } else { + e = firstErr(e, + b.wqkv.record(r, d.xn.b, d.qkv.b, k), + r.Copy2D(d.qkv.b, 0, 3*D, d.q.b, 0, D, k, D), + r.Copy2D(d.qkv.b, D, 3*D, d.k.b, 0, D, k, D), + r.Copy2D(d.qkv.b, 2*D, 3*D, d.v_.b, 0, D, k, D), + ) + } + } else { + e = firstErr(e, + b.wq.record(r, d.xn.b, d.q.b, k), + b.wk.record(r, d.xn.b, d.k.b, k), + b.wv.record(r, d.xn.b, d.v_.b, k), + ) + } + e = firstErr(e, + r.Blit(d.k.b, 0, b.kC, pos*D, k*D), + r.Blit(d.v_.b, 0, b.vC, pos*D, k*D), + r.MHA(d.q.b, b.kC, b.vC, d.attn.b, k, pos+k, D, H, H, dk, 1, 0, d.scale), + d.recordAttentionResidual(r, b, k), + r.LayerNorm(d.dx.b, b.g2, b.b2, d.xn2.b, k, D, d.eps), + b.w1.record(r, d.xn2.b, d.hid.b, k), + d.recordBiasGELU(r, b, k), + d.recordFFNResidual(r, b, k), + ) + if e != nil { + r.Free() + return nil, e + } + } + rows := k + if lastOnly { + // Only the final row's logits are wanted: move its residual to row 0 so the final LN + LM head + // project ONE row and the download shrinks from k·vocab to vocab (bit-identical — row 0 then + // holds exactly what row k-1 would have produced; a no-op when k==1). + if e := r.Blit(d.dx.b, (k-1)*D, d.dx.b, 0, D); e != nil { + r.Free() + return nil, e + } + rows = 1 + } + if e := firstErr( + r.LayerNorm(d.dx.b, d.lnfG.b, d.lnfB.b, d.xn.b, rows, D, d.eps), + d.head.record(r, d.xn.b, logits, rows), + r.Finish(), + ); e != nil { + r.Free() + return nil, e + } + r.Free() + out := make([]float32, rows*d.v) + if err := logits.DownloadF32(out); err != nil { + return nil, err + } + return out, nil +} + +// Generate prefills the whole prompt in ONE recorded StepN (§T423), then samples up to maxNew +// tokens (bounded by Ctx), returning prompt+generated ids. +func (d *GPTDecoder) Generate(prompt []int, maxNew int, s nlp.TokenSampler) ([]int, error) { + if len(prompt) == 0 { + return nil, fmt.Errorf("llamagpu(%s): gpt Generate needs a non-empty prompt", d.ops.name) + } + out := append([]int(nil), prompt...) + // generation samples only the post-prompt row, so project just the last prompt row's logits + // (StepNLast) instead of the full [len(prompt), vocab] LM head + download. + logits, err := d.StepNLast(prompt, 0) + if err != nil { + return nil, err + } + pos := len(prompt) + buf := make([]float64, d.v) + for range maxNew { + if pos >= d.maxLen { + break + } + for i, x := range logits { + buf[i] = float64(x) + } + next := s.SampleWithHistory(buf, out) + out = append(out, next) + l, err := d.Step(next, pos) + if err != nil { + return nil, err + } + logits = l + pos++ + } + return out, nil +} + +// Vocab returns the model's vocabulary size. +func (d *GPTDecoder) Vocab() int { return d.v } + +// Ctx returns the model's maximum context length (the KV-cache capacity). +func (d *GPTDecoder) Ctx() int { return d.maxLen } + +// Release frees all device buffers. +func (d *GPTDecoder) Release() { + d.fullLogits.release() + for _, b := range d.all { + if b != nil { + b.Release() + } + } + d.all = nil +} diff --git a/checks/testdata/ps6136-owner/before/decoder.go.txt b/checks/testdata/ps6136-owner/before/decoder.go.txt new file mode 100644 index 00000000..8ebface8 --- /dev/null +++ b/checks/testdata/ps6136-owner/before/decoder.go.txt @@ -0,0 +1,3896 @@ +// Package llamagpu decodes nlp.Llama models on the GPU with batched command buffers (ADR-0019): +// each per-token step records the whole layer stack into ONE command buffer over device-resident +// weights and KV cache, instead of paying a dispatch round-trip per op. Measured on a real model +// (D=512, GQA 8:2, 6 layers): 24× faster than nlp.Llama.DecodeStep on Metal and 21× on Vulkan, +// with token-for-token identical greedy output (§T404/§T409). +// +// Architecture coverage (CUDA): the same batched Decoder core drives 20 model families beyond Llama, +// each a New*CUDA constructor + greedy parity test, built from one composable set of gated +// generalizations (norm type/placement — pre / one-norm-parallel / two-norm-parallel / post-norm / +// sandwich; RMSNorm / LayerNorm±bias / (1+w); MLP — SwiGLU / GELU / squared-ReLU / GeGLU; position — +// full RoPE / partial RoPE / ALiBi; attention width — MHA / GQA / MQA / decoupled-head_dim; per-head +// & full-width QK-norm; attention soft-cap; config scalars; logit scale; √dim embed; biased +// projections/lm_head). Dense: StableLM, StarCoder2, Phi, GPT-NeoX, Qwen2, Qwen3, Phi-3, Granite, +// Cohere, Nemotron, Gemma, OLMo2, Falcon, Gemma2, MPT. Sparse Mixture-of-Experts (route → experts → +// combine, all in the pre-recorded command buffer via on-device routing): Mixtral, Qwen3-MoE, OLMoE, +// GraniteMoE, Qwen2-MoE (shared expert). These New*CUDA constructors are cuda-only where partial +// rotary / soft-cap / ALiBi / MoE kernels are involved; metal/vulkan return the corresponding stubs. +// +// Build [New] (Metal, darwin+cgo) or [NewVulkan] (vulkan build tag) from any *nlp.Llama — including +// one loaded via nlp.LlamaFromGGUF — then call [Decoder.Generate] with any nlp.TokenSampler +// (nlp.Sampler or nlp.Mirostat), or drive +// [Decoder.Step] / [Decoder.StepN] yourself. Beyond plain generation the package provides: +// +// - [Decoder.StepN]: a whole multi-token window in one recorded step. Generate uses it to prefill +// the prompt in a single dispatch round-trip — measured 41× faster than token-at-a-time (§T418). +// - [NewQuant] / [NewQuantF16KV] / [NewQuantVulkan]: decode a quantized model +// (nlp.QuantizeLlama or a quantized GGUF) with every projection held device-resident in its +// 4-8× smaller ggml form (§T415). NewQuantF16KV additionally halves retained Metal K/V cache +// storage while preserving the existing f32-cache constructor as the default. +// - [SpeculativeGenerate]: lossless speculative decoding — a small draft Decoder proposes, the +// target verifies the window in one StepN, and the output stays exactly target-distributed +// (§T419; ~2-3× at typical acceptance rates with a trained draft). +// - [PromptLookupGenerate]: draft-model-free speculative decoding — candidate continuations are +// copied from the sequence's own history by n-gram matching (§T426); lossless, and effective +// when the output repeats the input (summarization, RAG, code editing). +// - [MedusaGenerate]: Medusa chain drafting (§T446/§T447) — nlp.MedusaHeads trained on the base's +// own rollouts draft host-side for free, one batched StepN verifies the window. Measured 1.81× +// at 97% acceptance where draft-model speculative managed 1.12× (drafting cost is the +// difference on dispatch-bound decoders); greedy-anchored, not distribution-exact. +// +// The Decoder core is backend-agnostic over small buffer/recorder interfaces; the two adapters plug +// in the concrete backends, whose exported Recorder/DeviceBuffer APIs are identical by construction +// (§T391/§T408). +package llamagpu + +import ( + "errors" + "fmt" + "math" + + "github.com/jxsl13/goai/backend" + "github.com/jxsl13/goai/nlp" + "github.com/jxsl13/goai/nn" + "github.com/jxsl13/goai/tensor" +) + +// unarySiLU/binaryAdd/binaryMul are the shared kernel selectors (identical on both backends — +// they must match shaders/unary.comp / metal_bridge.m's unary switch and the binary op tables). +const ( + unarySiLU = 6 + unaryGELU = 9 + unaryReLU2 = 10 // squared ReLU (Nemotron); cuda-only + unarySigmoid = 11 // plain sigmoid (Qwen2-MoE shared-expert gate); cuda-only + unarySoftplus = 12 // softplus (Mamba Δ); cuda-only + unaryReLU = 13 // plain ReLU max(x,0) (T5 v1.0 FFN); cuda-only + binaryAdd = 0 + binaryMul = 2 + binarySwiGLU = 6 // fused silu(a)·b — one dispatch instead of SiLU+Mul (§T613) +) + +// buffer is device-resident storage (metal/cuda/vulkan DeviceBuffer). Decoder arithmetic buffers +// are f32. The opt-in Metal f16-KV capability uses this same lifetime-only interface for K/V cache +// storage; it never calls UploadF32/DownloadF32 on those half buffers. +type buffer interface { + UploadF32([]float32) error + DownloadF32([]float32) error + Release() +} + +// qweight is a backend-resident quantized weight (metal/vulkan *ResidentQWeight) — opaque to the +// core except for Close. +type qweight interface { + Close() error +} + +// binaryNRecorder is an OPTIONAL recorder capability (same shape as quantAccRecorder): an +// elementwise op bounded to the first n elements. Decoder buffers are sized for the whole context, +// so an unbounded add over one decoded row touches ctx x width elements instead of width. Recorders +// that do not implement it keep the unbounded call and stay correct, just slower. +type binaryNRecorder interface { + BinaryN(a, b, o buffer, op, n int) error +} + +// f16KVRecorder is the optional narrow capability needed by an f16 KV cache: convert f32 projection +// rows into retained binary16 cache storage, then read that cache in attention while Q/O and all +// accumulators remain f32. Keeping it outside recorder preserves every existing backend and default +// decoder path; only NewQuantF16KV supplies both this capability and an f16 cache allocator. +type f16KVRecorder interface { + Copy2DF32ToF16Pair( + kSrc buffer, kSrcOff, kSrcStride int, + vSrc buffer, vSrcOff, vSrcStride int, + kDst buffer, kDstOff, kDstStride int, + vDst buffer, vDstOff, vDstStride int, + rows, rowFloats int, + ) error + MHAF16KV(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error +} + +// ropeF16KVAppendRecorder is the optional Metal single-token capability that rotates contiguous +// Q/K projection rows while appending K/V directly to their binary16 cache row. Attention remains +// a separate dispatch because Metal provides no grid-wide barrier between the append and MHA work. +type ropeF16KVAppendRecorder interface { + RoPEF16KVAppend(q, k, v, inv, kCache, vCache buffer, + headsQ, headsK, hd, half, pos, cacheOff int, posDiv float32, + ) error +} + +// ropePairSplitRecorder is the Metal-only fused-QKV post-projection capability. It preserves the +// ordinary recorder interface for every other backend while allowing a grouped projection to +// rotate q/k and materialize contiguous q/k/v in one dispatch instead of rotate plus three copies. +type ropePairSplitRecorder interface { + RoPEPairSplit(qkv, inv, q, k, v buffer, + seq, stride, headsQ, offQ, headsK, offK, hd, half, vOff, vDim, pos int, posDiv float32, + ) error +} + +// ropePairF16KVAppendRecorder is the grouped-QKV counterpart of ropeF16KVAppendRecorder. It keeps +// Q/K in the combined projection row while appending its K/V bands directly into half caches. +type ropePairF16KVAppendRecorder interface { + RoPEPairF16KVAppend(qkv, inv, kCache, vCache buffer, + stride, headsQ, offQ, headsK, offK, hd, half, vOff, vDim, pos, cacheOff int, posDiv float32, + ) error +} + +// linear records o[m,N] = x[m,K]·W for one projection — either an f32 device weight (MatMul) or a +// resident quantized weight (QMatMulResident). This is what lets ONE Decoder core serve both the +// f32 and the quantized model (§T415). +type linear interface { + record(r recorder, x, o buffer, m int) error + // recordAdd records dst += x·W (the residual-add epilogue, §T613). The f32 weight fuses + // the add into the matmul (one dispatch); the quantized weight has no accumulate path, + // so it multiplies into scratch and adds — the pre-fusion two-dispatch shape. + recordAdd(r recorder, x, scratch, dst buffer, m, width int) error + // recordMoE is record() with per-token MoE expert gating: on a recorder that implements the gated + // path (CUDA Q8), a token whose routing weight moeGate[row*gateStride+ex]==0 skips the weight-stream + // (bit-identical to record + the RowAxpy weight-0 combine, ~E/K faster). Dense-falls-back otherwise. + recordMoE(r recorder, x, o buffer, m int, moeGate buffer, gateStride, ex int) error +} + +// moeGatedRecorder is the optional gated-expert-GEMV extension, implemented only by the CUDA recorder for +// Q8 experts; quantLinear.recordMoE type-asserts it and dense-falls-back when absent. +type moeGatedRecorder interface { + QMatMulResidentMoE(x buffer, w qweight, o buffer, m int, moeGate buffer, gateStride, ex int) error +} + +// quantAccRecorder is the optional residual-add-FUSED quant-GEMV extension: dst += x·dequant(w) in ONE +// kernel (beta=1), instead of GEMV→scratch (beta=0) plus a separate Binary add. Implemented by the CUDA +// recorder for the Q8 decode GEMV; quantLinear.recordAdd type-asserts it for m==1 and falls back to the +// two-op shape when absent or for a format without a fused path (errQuantAccUnsupported). +type quantAccRecorder interface { + QMatMulResidentAcc(x buffer, w qweight, dst buffer, m int) error +} + +// dependencyBarrierRecorder is implemented by the Metal adapter. Ordinary/profiling Metal +// recorders make Barrier a no-op; the decode-only concurrent recorder emits a buffer-scope fence. +// Other backends preserve their established command ordering without implementing this extension. +type dependencyBarrierRecorder interface { + Barrier() error +} + +func recordBarrier(r recorder) error { + if br, ok := r.(dependencyBarrierRecorder); ok { + return br.Barrier() + } + return nil +} + +// errQuantAccUnsupported signals a weight format with no fused residual-add GEMV — recordAdd then uses +// the GEMV-into-scratch + Binary-add path. +var errQuantAccUnsupported = errors.New("llamagpu: no fused quant residual-add for this weight format") + +// swiGLUHalvesRecorder is the optional column-fused gate|up SwiGLU extension: given a [rows, 2·hidden] +// gate|up GEMV output, compute out[r,i] = silu(gu[r,i])·gu[r,hidden+i] in ONE op — the launch-count +// twin of the residual fusion for the FFN. Implemented by the CUDA recorder; recordFFN uses it when a +// block carries the fused wGU weight (built only when ops.fusedGateUp), else the two-GEMV path. +type swiGLUHalvesRecorder interface { + SwiGLUHalves(gu, out buffer, rows, hidden int) error +} + +// geGLUHalvesRecorder is the GeGLU (Gemma) twin: out[r,i] = GELU(gu[r,i])·gu[r,hidden+i] over a +// column-fused gate|up buffer, one op replacing two GEMVs + Unary GELU + Binary mul. +type geGLUHalvesRecorder interface { + GeGLUHalves(gu, out buffer, rows, hidden int) error +} + +type f32Linear struct { + w buffer + k, n int +} + +func (l f32Linear) record(r recorder, x, o buffer, m int) error { + return r.MatMul(x, l.w, o, m, l.k, l.n) +} + +func (l f32Linear) recordAdd(r recorder, x, _, dst buffer, m, _ int) error { + return r.MatMulAcc(x, l.w, dst, m, l.k, l.n) +} + +type quantLinear struct{ w qweight } + +func (l quantLinear) record(r recorder, x, o buffer, m int) error { + return r.QMatMulResident(x, l.w, o, m) +} + +func (l quantLinear) recordAdd(r recorder, x, scratch, dst buffer, m, width int) error { + // Decode / small batch (m<6): fuse the residual add into the quant kernel (beta=1) — ONE launch into + // dst instead of QMatMulResident(beta=0)→scratch + a separate Binary add, saving the add launch and + // the scratch HBM round-trip. The old code fused only m==1 because a beta=1 GEMV re-streams the weight + // M× per output (undoing the weight-read-once routing); that objection is void now that the Acc paths + // route small batches to the SAME weight-read-once MT kernel their plain twins use (Q8 #1011/#1012, + // K-quants #1013), so beta=1 does strictly less work than beta=0+add. Gate m<6: below it every quant + // uses the MT/GEMV for BOTH plain and Acc, so fused is guaranteed ≤ split. At m>=6 the plain path may + // switch to a faster kernel the Acc path lacks (Q8→int8-MMQ at m>=6, K-quants→WMMA at m>=48), so the + // split — which gets that faster kernel for the GEMM — is kept. Bit-identical either way (the Acc MT is + // bit-identical to plain MT, and beta=1 folds the add exactly; TestCUDAKQuantAccMTMatchesGEMV @m=4). + if m < 6 { + if acc, ok := r.(quantAccRecorder); ok { + if err := acc.QMatMulResidentAcc(x, l.w, dst, m); err != errQuantAccUnsupported { + return err + } + } + } + // Bound the add to the m rows actually written. Without this it runs over the whole + // context-sized buffer — 1024x the traffic at ctx=1024, and it dominated decode. + if bn, ok := r.(binaryNRecorder); ok && width > 0 { + e := r.QMatMulResident(x, l.w, scratch, m) + e = firstErr(e, recordBarrier(r)) + return firstErr(e, bn.BinaryN(dst, scratch, dst, binaryAdd, m*width)) + } + e := r.QMatMulResident(x, l.w, scratch, m) + e = firstErr(e, recordBarrier(r)) + return firstErr(e, r.Binary(dst, scratch, dst, binaryAdd)) +} + +func (l f32Linear) recordMoE(r recorder, x, o buffer, m int, _ buffer, _, _ int) error { + return l.record(r, x, o, m) // f32 MoE experts (rare) use the dense path +} + +func (l quantLinear) recordMoE(r recorder, x, o buffer, m int, moeGate buffer, gateStride, ex int) error { + if mg, ok := r.(moeGatedRecorder); ok { + return mg.QMatMulResidentMoE(x, l.w, o, m, moeGate, gateStride, ex) + } + return l.record(r, x, o, m) +} + +// recorder is one open batched command buffer (metal.Recorder / vulkan.Recorder); the adapter +// asserts the buffer args back to its concrete type. +type recorder interface { + RMSNorm(x, g, o buffer, rows, dim int, eps float32) error + LayerNorm(x, g, b, o buffer, rows, dim int, eps float32) error + AddBias(x, b, o buffer, rows, n int) error + MatMul(a, b, c buffer, m, k, n int) error + // MatMulAcc records c += a·b — the residual-add epilogue (§T613): the projection lands + // directly in the running residual stream, saving the separate elementwise-add dispatch. + MatMulAcc(a, b, c buffer, m, k, n int) error + RoPE(q, inv, o buffer, seq, width, heads, hd, half, pos int, posDiv float32) error + // RoPEAt rotates a sub-row view living at float-element offset `off` inside a wider + // buffer — the q/k bands of a fused QKV projection output (§T613). width acts as the + // ROW STRIDE; only the first heads·hd columns of each row are rotated, in place. + RoPEAt(q, inv, o buffer, off, seq, width, heads, hd, half, pos int, posDiv float32) error + // RoPEPair rotates the q band (headsQ heads at offQ) AND the k band (headsK at offK) of + // a fused QKV buffer in ONE dispatch (§T613), rows `stride` floats wide, in place. + RoPEPair(qkv, inv buffer, seq, stride, headsQ, offQ, headsK, offK, hd, half, pos int, posDiv float32) error + // RoPEPartialPair is the partial-rotary sibling of RoPEPair: only the first rotaryDim + // channels of each head in the q and k bands rotate (GPT-NeoX/Phi/StableLM partial + // rotary). inv must be the [rotaryDim/2] frequency table. Implemented on cuda; metal and + // vulkan return an unsupported error until a partial-rotary decoder targets them. + RoPEPartialPair(qkv, inv buffer, seq, stride, headsQ, offQ, headsK, offK, hd, rotaryDim, pos int, posDiv float32) error + Blit(src buffer, srcOff int, dst buffer, dstOff, n int) error + // Copy2D moves a strided rows×rowFloats sub-matrix (the fused-QKV band extraction, §T613): + // row r copies from src[srcOff+r·srcStride:] to dst[dstOff+r·dstStride:]. + Copy2D(src buffer, srcOff, srcStride int, dst buffer, dstOff, dstStride, rows, rowFloats int) error + MHA(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error + // MHACap is MHA with a Gemma-2 attention-logit soft-cap (cap·tanh(score·scale/cap) before the + // mask+softmax). Implemented on cuda; metal/vulkan return unsupported until a softcap decoder + // targets them (the RoPEPartialPair pattern). + MHACap(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale, cap float32) error + // MHAALiBi is MHA with an ALiBi position bias (MPT): slopeₕ·(j−qabs) added to each scaled score + // before the mask+softmax; slopes holds `heads` per-head slopes. cuda-only (metal/vulkan stub). + MHAALiBi(q, k, v, o, slopes buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error + // MHABias is MHA with a PRECOMPUTED per-head additive bias [heads,sq,sk] added to the scaled scores + // before softmax — the T5 relative-position bias (a learned table, not ALiBi's slope). No RoPE; + // causal==0 unmasks all keys (T5's bidirectional encoder). cuda-only (metal/vulkan stub). + MHABias(q, k, v, o, bias buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error + // MoEGate writes Mixtral top-k renormalized routing weights [rows·e] from router logits [rows·e]. + // RowAxpy accumulates dst[r,:] += arow[r]·src[r,:] (the MoE weighted combine). Both cuda-only. + MoEGate(logits, weights buffer, rows, e, k, raw int, scale float32) error + RowAxpy(dst, src, arow buffer, rows, cols int) error + // MHARect is multi-head attention with a rectangular head shape: query/key share head width dqk, + // value has a different head width dv (o is heads·dv wide). The DeepSeek-V2 MLA shape. cuda-only. + MHARect(q, k, v, o buffer, sq, sk, heads, kvHeads, dqk, dv, causal, window int, scale float32) error + // SSMStep advances ONE Mamba selective-scan timestep (decode): given u/Δ/A/B/C[d,·] and the skip + // D[d], it updates the [d,n] recurrent state h IN PLACE and writes y[d]. Conv1DStep advances one + // causal depthwise-conv step: out[c]=Σ_k w[c,k]·window[c,k]+b[c], then shifts the [d,k-1] state and + // appends x. Both carry per-block state across calls (no KV cache) — the GPU Mamba block. cuda-only. + SSMStep(u, delta, a, b, c, dskip, h, y buffer, d, n int) error + Conv1DStep(x, w, b, state, out buffer, d, k int) error + // SSDStep advances ONE Mamba-2 SSD timestep: H heads with a SCALAR per-head decay a=exp(Δ[h]·A[h]) + // and B/C shared across a group (g=h/(H/G)); updates the per-head state[H,N,P] in place and writes + // y[H·P]. x/y are [H·headDim]=intermediate; delta/a/dskip are [H]; b/c are [G·N]. cuda-only. + SSDStep(x, delta, a, b, c, dskip, state, y buffer, heads, headDim, groups, n int) error + // WKVStep advances ONE RWKV-4 WKV recurrence timestep: k/v/w/u/out are [d] (w = per-channel decay + // exp(WLog), u = current-token bonus); the running state aa/bb/pp [d] is updated in place (fresh + // sequence: aa=bb=0, pp=−1e38). No KV cache — RWKV's O(1) recurrent inference. cuda-only. + WKVStep(k, v, w, u, aa, bb, pp, out buffer, d int) error + Unary(x, o buffer, op int) error + Binary(a, b, o buffer, op int) error + QMatMulResident(x buffer, w qweight, o buffer, m int) error + // Commit submits without waiting; Wait blocks until done (§T614 encode-overlap split). + // Backends without async submit implement Commit as a deferral and Wait as Finish. + Commit() error + Wait() error + Finish() error + Free() +} + +// backendOps is the per-backend constructor vtable the adapters fill in. +type backendOps struct { + name string + newBuffer func([]float32) (buffer, error) + newF16KVBuffer func(int) (buffer, error) // nil unless this constructor opts into binary16 K/V storage + newRecorder func() (recorder, error) + // newDecodeRecorder optionally specializes only the single-token Step graph. Prefill and every + // other recorder consumer retain newRecorder; nil falls back to newRecorder. + newDecodeRecorder func() (recorder, error) + uploadQWeight func(weight []byte, qt uint32, n, k int) (qweight, error) // resident quantized upload + groupQWeights func(weights []qweight) (qweight, error) // optional heterogeneous prefill group + // quantizeF32 quantizes a host f32 weight [in,out] to a resident Q8_0 device weight in one call + // (no pre-quantized ggml bytes needed). Set only by the cuda Q8 recurrent-decode constructors — + // nil elsewhere, so the caller keeps the f32 path. Lets a decoder trade the f32 weight's bandwidth + // (the bottleneck for the weight-bound recurrent decode) for a ~4× smaller Q8 GEMV. + quantizeF32 func(w *tensor.Tensor) (qweight, error) + // asyncEncode: the backend can encode a second recorder while one executes (§T614). + // Metal command buffers are independent objects → true; the vulkan bridge keeps ONE + // global recording context → false (pre-encoding there would clobber the in-flight one). + asyncEncode bool + // fusedGateUp: the backend's recorder implements swiGLUHalvesRecorder, so a SwiGLU FFN's + // ffn_gate|ffn_up can be column-fused into ONE GEMV (into the gu buffer) + one SwiGLUHalves op, + // instead of two GEMVs + a Binary SwiGLU. Set by the CUDA SwiGLU constructors; nil elsewhere. + fusedGateUp bool + // fusedF32QKV opts a backend into GPT's single-resident-weight QKV path. The recorder must + // implement f32QKVBandsRecorder so large prefills can address each output band without copies. + fusedF32QKV bool + // fusedBiasGELU opts GPT into the recorder's bounded bias-plus-exact-GELU epilogue. Recorders + // without the optional capability retain the portable AddBias plus Unary chain. + fusedBiasGELU bool +} + +// moeFFN is one sparse-MoE expert: a SwiGLU FFN (gate/up/down) with its own weights. +type moeFFN struct { + wG, wU, wD linear + wGU linear // column-fused ffn_gate|ffn_up (built only when ops.fusedGateUp); nil ⇒ wG/wU path +} + +type block struct { + wq, wk, wv, wo, wG, wU, wD linear + // wqkv is the fused QKV projection (§T613): one [in, D+2·kvDim] weight whose output row + // is [q | k | v]. Non-nil = Step/StepN run the fused chain (one matmul instead of three); + // nil = fall back to wq/wk/wv (quantized blocks whose projections mix quant types). + wqkv linear + // wqkvPrefill is the Metal-only mixed-quant companion: StepN uses one exact post-expansion + // projection, while Step keeps wq/wk/wv and therefore preserves the established decode kernels. + wqkvPrefill linear + // wGU is the column-fused ffn_gate|ffn_up projection: one [in, 2·hidden] weight whose output row + // is [gate | up]. Non-nil (built only when ops.fusedGateUp, i.e. CUDA SwiGLU) = recordFFN runs one + // GEMV into the gu buffer + one SwiGLUHalves; nil = the separate wG/wU two-GEMV SwiGLU path. + wGU linear + gAttn, gFFN, kC, vC buffer + bAttn, bFFN buffer // LayerNorm betas (nil ⇒ RMSNorm, the Llama default) + // projection biases (nil ⇒ no bias): the fused-QKV bias [D+2·kvDim], the o-proj bias [D], + // and the GELU-MLP fc [hidden] / proj [D] biases — GPT-NeoX/Phi/StarCoder2 carry them. + qkvBias, oBias, fcBias, projBias buffer + qN, kN buffer // per-head query/key RMSNorm gains [dk] (Qwen3); nil ⇒ no QK-norm + gAttn2, gFFN2 buffer // sandwich-norm OUTPUT gains (Gemma2 post_attention/post_feedforward); nil otherwise + moeRouter linear // sparse-MoE gate: dim → nExperts logits (Mixtral); nil unless d.moe + moeExperts []moeFFN // the per-expert SwiGLU FFNs (Mixtral); len == d.nExperts + moeShared moeFFN // Qwen2-MoE shared expert (SwiGLU on every token); zero-value when absent + moeSharedGate linear // Qwen2-MoE shared-expert gate: dim → 1 logit, sigmoid-scaled; nil when absent + wqA, wqB, wkvA, wkvB linear // DeepSeek-V2 MLA low-rank projections (q_a/q_b/kv_a/kv_b); nil unless d.mla + gQA, gKvA buffer // MLA latent RMSNorm gains (q_a_layernorm / kv_a_layernorm) + mInX, mInZ, mDtLow, mDtProj, mBProj, mCProj, mOutProj linear // Mamba block projections; nil unless d.mamba + mConvW, mConvB, mA, mDskip, mDtBias buffer // Mamba conv filters+bias, A=−exp(A_log) [dInner,N], D skip [dInner], Δ bias [dInner] + mConvState, mSsmState buffer // per-block Mamba decode state: conv window [dInner,dConv-1], ssm state [dInner,N] + mDtNorm, mBNorm, mCNorm buffer // Jamba mixer RMSNorm gains on Δ_low [dtRank] / B [N] / C [N]; nil for plain Mamba (recordMambaBlock applies them only when set) + jMamba bool // Jamba: this layer's mixer is Mamba (else NoPE attention); only meaningful when d.jamba + m2InProj, m2OutProj linear // Mamba-2 fused in_proj [d,projSize] / out_proj [inter,d]; nil unless d.mamba2 + m2ConvW, m2ConvB, m2A, m2DtBias, m2Dskip, m2NormW buffer // Mamba-2 conv filters+bias [convDim,·], A/Δbias/D [numHeads], gated-RMSNorm gain [inter] + m2ConvState, m2SsdState buffer // Mamba-2 state: conv window [convDim,dConv-1], SSD state [numHeads·N·headDim] + // RWKV-4 (nil unless d.rwkv): time-mix reuses wq/wk/wv/wo = Wr/Wk/Wv/Wo, channel-mix reuses + // wG/wD = CWk/CWv; LN1 = gAttn/bAttn, LN2 = gFFN/bFFN. rwCWr is the channel-mix receptance. + rwCWr linear // channel-mix receptance CWr [Dim,Dim] + rwMuR, rwMuK, rwMuV, rwCMuR, rwCMuK buffer // token-shift interpolators μ [Dim] + rwOmR, rwOmK, rwOmV, rwOmCR, rwOmCK buffer // (1−μ) precomputed [Dim] + rwW, rwU buffer // WKV per-channel decay w=exp(WLog) [Dim], bonus u [Dim] + rwPrevTM, rwPrevCM, rwAA, rwBB, rwPP buffer // per-block RWKV recurrent state [Dim] (token-shift + WKV aa/bb/pp) +} + +// Decoder holds a Llama's weights + KV cache as device-resident buffers and runs one batched decode +// step per token. Not safe for concurrent use — one Decoder per goroutine. Release when done. +type Decoder struct { + ops backendOps + d, h, kvH, dk, kvDim, half, hidden, v, maxLen int + qDim int // attention Q/O width = h·dk; == d (model dim) unless head_dim is decoupled (Gemma) + rotaryDim int // rotated channels/head; == dk for full RoPE (Llama), < dk for partial rotary (GPT-NeoX/Phi/StableLM) + lnBias bool // true ⇒ LayerNorm-with-bias norms (StableLM/Phi/StarCoder2); false ⇒ RMSNorm (Llama) + ffnGELU bool // true ⇒ 2-layer GELU MLP (fc→GELU→proj: GPT-NeoX/Phi/StarCoder2); false ⇒ SwiGLU (Llama/StableLM) + ffnReLU2 bool // true ⇒ 2-layer squared-ReLU MLP (fc→relu²→proj: Nemotron); mutually exclusive with ffnGELU + ffnGEGLU bool // true ⇒ GeGLU MLP: down(GELU(gate)⊙up) — Gemma's GELU-gated variant of SwiGLU + parallelRes bool // true ⇒ parallel residual: attn+FFN both read norm(x0), sum onto x (GPT-NeoX/Phi/Cohere); false ⇒ sequential (Llama) + parallelTwoNorm bool // parallel residual with SEPARATE attn/FFN norms both over x0 (GPT-NeoX); needs the FFN norm computed pre-attn-add + qkNorm bool // true ⇒ RMSNorm on Q and K before RoPE (Qwen3 per-head, OLMo2 full-width); false otherwise + qkNormFull bool // with qkNorm: true ⇒ one RMSNorm over the whole q/k projection (OLMo2), false ⇒ per-head (Qwen3) + postNorm bool // true ⇒ post-norm blocks (OLMo2): sublayer reads the raw residual, its OUTPUT is normed before the add + sandwich bool // true ⇒ sandwich norms (Gemma2): pre-norm the input (gAttn/gFFN) AND post-norm the output (gAttn2/gFFN2) + attnCap float32 // >0 ⇒ Gemma2 attention-logit soft-cap magnitude (via MHACap); 0 ⇒ plain MHA + noRope bool // true ⇒ no rotary position embedding (MPT: position enters via ALiBi only) + aliBiSlopes buffer // non-nil ⇒ per-head ALiBi slopes [heads]; attention runs via MHAALiBi (MPT) + moe bool // true ⇒ the FFN sublayer is a sparse Mixture-of-Experts (Mixtral): route → experts → combine + nExperts, topK int // MoE expert count and experts-per-token (when moe) + moeRaw bool // true ⇒ DeepSeek-V2 raw-softmax·routedScale gating (no top-k renormalization); false ⇒ Mixtral renorm + routedScale float32 // DeepSeek-V2 routed-expert weight multiplier (routed_scaling_factor); 1 otherwise + mla bool // true ⇒ DeepSeek-V2 Multi-head Latent Attention (rectangular, low-rank latent KV, decoupled RoPE) + mamba bool // true ⇒ Mamba selective-state-space blocks (no attention, no KV cache; linear-time recurrent decode) + dInner, mambaN, dConv, dtRank int // Mamba block dims: inner width (e·d), SSM state size, conv kernel, Δ low-rank + mamba2 bool // true ⇒ Mamba-2 SSD blocks (scalar-per-head decay, grouped B/C, fused in_proj, gated RMSNorm) + m2H, m2P, m2G, m2N, m2Conv, m2Inter, m2CD int // Mamba-2 dims: num_heads, head_dim, n_groups, state_size, conv kernel, intermediate, conv_dim + rwkv bool // true ⇒ RWKV-4 blocks (WKV time-mix + gated squared-ReLU channel-mix; no attention/KV cache; O(1) recurrent) + rwHidden int // RWKV channel-mix hidden width (CWk: Dim→Hidden, CWv: Hidden→Dim) + jamba bool // true ⇒ Jamba hybrid stack: each layer is Mamba-mixer (jMamba) OR NoPE GQA attention, then a per-block MoE/dense FFN; sequential like Mamba + qkNope, qkRope, vHead, qkHead int // MLA per-head dims: nope/rope query-key parts, value width, qkHead=qkNope+qkRope + qLoRA, kvLoRA int // MLA query/kv latent compression ranks (q_lora_rank / kv_lora_rank) + mlaScale, mlaPosDiv float32 // MLA pre-softmax scale (1/√qkHead default) and the QKRope rope posDiv + eps, posDiv, scale float32 + embMult float32 // gathered embeddings ×= embMult (IBM Granite EmbeddingMult); 1 for everything else + f16KV bool // opt-in: K/V cache storage is IEEE binary16; compute remains f32 + + blocks []block + out linear + gFinal, bFinal, outBias, dinv *bufSlot + dx, xn, xn2, q, k, v_, attn, ao, gate, up, mo, logits *bufSlot + gu *bufSlot // fused gate|up GEMV output [·, 2·hidden] (nil unless ops.fusedGateUp) + moeGate, moeW, moeCol *bufSlot // sparse-MoE scratch: router logits [·,E], routing weights [·,E], one weight column [·] + mlaCQ, mlaQ, mlaComp, mlaLatent, mlaKV, mlaAttn, mlaInv *bufSlot // MLA scratch: compressed query, fused Q, kv_a out, normed kv latent, kv_b out, attn out, QKRope freqs + mbXin, mbZ, mbXc, mbDtLow, mbDelta, mbB, mbC, mbY *bufSlot // Mamba scratch (rows=1): in_x, gate branch, conv+SiLU out, Δ_low, Δ, B, C, scan out + m2Proj, m2Z, m2XBC, m2Dt, m2Xc, m2X, m2B, m2C, m2Y *bufSlot // Mamba-2 scratch (rows=1): in_proj out, gate z, conv input xBC, Δ pre-act, conv+SiLU out, value x, B, C, SSD out + rwPreLNg, rwPreLNb *bufSlot // RWKV ln0 (pre-block-0 LayerNorm) γ/β + rwMix, rwR, rwK, rwV, rwWkv, rwO, rwT1, rwT2, rwCr, rwCk, rwCv *bufSlot // RWKV scratch (rows=1): token-shift mix, r/k/v, WKV out, time-mix out, 2 mix temps, chan-mix cr, ck [Hidden], cv + qkv *bufSlot // fused QKV output rows [·, d+2·kvDim] (§T613) + pending recorder // pre-encoded next-step command buffer (§T614) + pendingPos int // position pending encodes; -1 = none + table *tensor.Tensor // token embedding, host-side gather source + invHost []float32 // RoPE inverse freqs (uploaded by allocScratch) + all []buffer + qweights []qweight // resident quantized weights (quant decoder) +} + +// bufSlot wraps a buffer so the struct fields read naturally while sharing the release list. +type bufSlot struct{ b buffer } + +// newDecoderCommon sets the geometry, RoPE frequencies and scratch buffers shared by the f32 and +// quantized constructors. +func newDecoderCommon(cfg nlp.LlamaConfig, tokEmb *tensor.Tensor, ops backendOps) (*Decoder, error) { + d := &Decoder{ + ops: ops, + d: cfg.Dim, h: cfg.Heads, kvH: cfg.KVHeads, hidden: cfg.Hidden, v: cfg.Vocab, + maxLen: cfg.Ctx, eps: float32(cfg.Eps), table: tokEmb, + pendingPos: -1, embMult: 1, + } + if d.kvH <= 0 { + d.kvH = d.h + } + if d.h <= 0 || d.d%d.h != 0 { + return nil, fmt.Errorf("llamagpu(%s): dim %d not divisible by heads %d", ops.name, d.d, d.h) + } + d.dk = d.d / d.h + d.qDim = d.h * d.dk // = d.d for standard models; decoupled-head_dim constructors (Gemma) override + d.kvDim = d.kvH * d.dk + d.half = d.dk / 2 + d.rotaryDim = d.dk // full RoPE by default; partial-rotary constructors override before allocScratch + d.scale = float32(1.0 / math.Sqrt(float64(d.dk))) + + invF64, posDiv64 := backend.RoPEFreqs(d.dk, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, d.half) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + return d, nil +} + +// ropeQK records RoPE over the q and k bands of the fused QKV buffer in one pass: full rotary +// (RoPEPair) when rotaryDim == dk (Llama), else partial rotary (RoPEPartialPair — only the first +// rotaryDim channels of each head rotate; GPT-NeoX/Phi/StableLM). inv (d.dinv) is sized to +// rotaryDim/2 so the same buffer feeds both. +// recordFFN records the feed-forward sublayer with the residual-add epilogue (dx += ffn(xn2)): +// a 2-layer GELU MLP (fc → GELU → proj, GPT-NeoX/Phi/StarCoder2) when ffnGELU, else the SwiGLU +// (silu(gate)·up → down, Llama/StableLM). wG/wU/wD are c_fc/·/c_proj or gate/up/down accordingly. +// recordQKVProj records the fused QKV projection (xn·Wqkv) plus its optional bias — the bias is +// added BEFORE RoPE, per the [q|k|v] band layout (GPT-NeoX/Phi/StarCoder2 have biased q/k/v). +// recordAttnNorm produces the attention sublayer's input in d.xn. Pre-norm (Llama & most): xn = +// norm(x0) (plus norm2(x0)→xn2 for two-norm parallel). Post-norm (OLMo2): attention reads the RAW +// residual, so xn is just a copy of x0 and the norm happens later on the attention output. +func (d *Decoder) recordAttnNorm(r recorder, b block, rows int) error { + if d.postNorm { + return r.Blit(d.dx.b, 0, d.xn.b, 0, rows*d.d) + } + e := d.norm(r, d.dx.b, b.gAttn, b.bAttn, d.xn.b, rows) + if d.parallelTwoNorm { // norm2(x0) BEFORE the attn add, for the FFN branch + e = firstErr(e, d.norm(r, d.dx.b, b.gFFN, b.bFFN, d.xn2.b, rows)) + } + return e +} + +func (d *Decoder) recordQKVProj(r recorder, b block, rows int) error { + return d.recordQKVProjWith(r, b.wqkv, b.qkvBias, rows) +} + +func (d *Decoder) recordQKVProjWith(r recorder, w linear, bias buffer, rows int) error { + e := w.record(r, d.xn.b, d.qkv.b, rows) + if bias != nil { + e = firstErr(e, r.AddBias(d.qkv.b, bias, d.qkv.b, rows, d.qDim+2*d.kvDim)) + } + return e +} + +// recordQKNorm records Qwen3's per-head RMSNorm on the Q and K bands of the fused QKV buffer, +// applied BEFORE RoPE. Each head's dk channels are normed independently with a shared gain — the +// reference reshapes [seq, heads·dk] → [seq·heads, dk], RMSNorms, reshapes back. The heads of one +// token are contiguous but tokens are stride-separated inside qkv, so we Copy2D the band out to a +// packed [seq, band] scratch (reusing d.q / d.k, unused in the fused path until after RoPE), +// RMSNorm it as seq·heads rows of dk, and Copy2D it back. No-op when qkNorm is unset. +func (d *Decoder) recordQKNorm(r recorder, b block, seq, stride int) error { + if !d.qkNorm { + return nil + } + // rows×dim for the RMSNorm: OLMo2 norms the WHOLE q/k projection as one vector (rows=seq, + // dim=qDim/kvDim, gain [qDim]/[kvDim]); Qwen3 norms each head (rows=seq·heads, dim=dk, gain [dk]). + qRows, qCols, kRows, kCols := seq*d.h, d.dk, seq*d.kvH, d.dk + if d.qkNormFull { + qRows, qCols, kRows, kCols = seq, d.qDim, seq, d.kvDim + } + e := r.Copy2D(d.qkv.b, 0, stride, d.q.b, 0, d.qDim, seq, d.qDim) // extract Q band [seq, qDim] + e = firstErr(e, r.RMSNorm(d.q.b, b.qN, d.q.b, qRows, qCols, d.eps)) + e = firstErr(e, r.Copy2D(d.q.b, 0, d.qDim, d.qkv.b, 0, stride, seq, d.qDim)) // write Q band back + e = firstErr(e, r.Copy2D(d.qkv.b, d.qDim, stride, d.k.b, 0, d.kvDim, seq, d.kvDim)) + e = firstErr(e, r.RMSNorm(d.k.b, b.kN, d.k.b, kRows, kCols, d.eps)) + e = firstErr(e, r.Copy2D(d.k.b, 0, d.kvDim, d.qkv.b, d.qDim, stride, seq, d.kvDim)) + return e +} + +// recordOProj records the attention output projection with its residual-add epilogue (dx += +// attn·Wo) plus the optional o-bias broadcast onto the residual stream. +// recordMHA runs attention against the Decoder's configured cache storage. The f16 capability is +// selected only by the opt-in constructor; all established constructors still take the exact f32 +// calls below. Initial batched prefill can explicitly pass f16=false to recordMHAStorage so FlashMM +// consumes its already-resident f32 K/V rows while those rows are also converted into the cache. +func (d *Decoder) recordMHA(r recorder, q, kC, vC, o buffer, sq, sk int) error { + return d.recordMHAStorage(r, q, kC, vC, o, sq, sk, d.f16KV) +} + +func (d *Decoder) recordMHAStorage(r recorder, q, kC, vC, o buffer, sq, sk int, f16 bool) error { + if d.aliBiSlopes != nil { + return r.MHAALiBi(q, kC, vC, o, d.aliBiSlopes, sq, sk, d.qDim, d.h, d.kvH, d.dk, 1, 0, d.scale) + } + if d.attnCap > 0 { + return r.MHACap(q, kC, vC, o, sq, sk, d.qDim, d.h, d.kvH, d.dk, 1, 0, d.scale, d.attnCap) + } + if f16 { + fr, ok := r.(f16KVRecorder) + if !ok { + return fmt.Errorf("llamagpu(%s): f16 KV cache requires recorder conversion/attention capability", d.ops.name) + } + return fr.MHAF16KV(q, kC, vC, o, sq, sk, d.qDim, d.h, d.kvH, d.dk, 1, 0, d.scale) + } + return r.MHA(q, kC, vC, o, sq, sk, d.qDim, d.h, d.kvH, d.dk, 1, 0, d.scale) +} + +func (d *Decoder) recordKVPairBlit(r recorder, + kSrc buffer, kSrcOff int, vSrc buffer, vSrcOff int, + kDst, vDst buffer, dstOff, n int, +) error { + if !d.f16KV { + return firstErr( + r.Blit(kSrc, kSrcOff, kDst, dstOff, n), + r.Blit(vSrc, vSrcOff, vDst, dstOff, n), + ) + } + fr, ok := r.(f16KVRecorder) + if !ok { + return fmt.Errorf("llamagpu(%s): f16 KV cache requires recorder conversion capability", d.ops.name) + } + return fr.Copy2DF32ToF16Pair( + kSrc, kSrcOff, n, + vSrc, vSrcOff, n, + kDst, dstOff, n, + vDst, dstOff, n, + 1, n, + ) +} + +func (d *Decoder) recordKVPairCopy2D(r recorder, + kSrc buffer, kSrcOff, kSrcStride int, + vSrc buffer, vSrcOff, vSrcStride int, + kDst, vDst buffer, dstOff, dstStride, rows, rowFloats int, +) error { + if !d.f16KV { + return firstErr( + r.Copy2D(kSrc, kSrcOff, kSrcStride, kDst, dstOff, dstStride, rows, rowFloats), + r.Copy2D(vSrc, vSrcOff, vSrcStride, vDst, dstOff, dstStride, rows, rowFloats), + ) + } + fr, ok := r.(f16KVRecorder) + if !ok { + return fmt.Errorf("llamagpu(%s): f16 KV cache requires recorder conversion capability", d.ops.name) + } + return fr.Copy2DF32ToF16Pair( + kSrc, kSrcOff, kSrcStride, + vSrc, vSrcOff, vSrcStride, + kDst, dstOff, dstStride, + vDst, dstOff, dstStride, + rows, rowFloats, + ) +} + +func (d *Decoder) recordOProj(r recorder, b block, rows int) error { + if d.postNorm || d.sandwich { + // Output-normed residual: ao = attn·Wo → norm(ao) → dx += ao. OLMo2 post-norm reuses gAttn; + // Gemma2 sandwich has a SEPARATE post_attention_layernorm gain (gAttn2), the input already + // having been pre-normed with gAttn. + g, bta := b.gAttn, b.bAttn + if d.sandwich { + g, bta = b.gAttn2, nil + } + return firstErr( + b.wo.record(r, d.attn.b, d.ao.b, rows), + d.norm(r, d.ao.b, g, bta, d.ao.b, rows), + d.binElem(r, d.dx.b, d.ao.b, d.dx.b, binaryAdd, rows, d.d), + ) + } + e := b.wo.recordAdd(r, d.attn.b, d.ao.b, d.dx.b, rows, d.d) + if b.oBias != nil { + e = firstErr(e, r.AddBias(d.dx.b, b.oBias, d.dx.b, rows, d.d)) + } + return e +} + +// recordDownProj records the FFN down-projection's residual epilogue: dx += act·Wdown, either fused +// (pre-norm) or, for OLMo2 post-norm, as mo = act·Wdown → post_feedforward_layernorm(mo) → dx += mo. +func (d *Decoder) recordDownProj(r recorder, b block, rows int) error { + if d.postNorm || d.sandwich { + g, bta := b.gFFN, b.bFFN + if d.sandwich { + g, bta = b.gFFN2, nil // Gemma2 post_feedforward_layernorm (input pre-normed with gFFN) + } + return firstErr( + b.wD.record(r, d.gate.b, d.mo.b, rows), + d.norm(r, d.mo.b, g, bta, d.mo.b, rows), + d.binElem(r, d.dx.b, d.mo.b, d.dx.b, binaryAdd, rows, d.d), + ) + } + return b.wD.recordAdd(r, d.gate.b, d.mo.b, d.dx.b, rows, d.d) +} + +// recordLogits records the final projection to logits (xn·Out) plus the optional output bias +// (Phi's biased untied lm_head; nil ⇒ no bias, as for Llama/StableLM/StarCoder2). +func (d *Decoder) recordLogits(r recorder, rows int) error { + e := d.out.record(r, d.xn.b, d.logits.b, rows) + if d.outBias != nil { + e = firstErr(e, r.AddBias(d.logits.b, d.outBias.b, d.logits.b, rows, d.v)) + } + return e +} + +// recordFFNSublayer records the FFN norm + the FFN with its residual-add epilogue. Sequential +// (Llama/StableLM/StarCoder2): norm the post-attention residual, then FFN(that). PARALLEL residual +// (Phi/Cohere one-norm): the FFN reuses the SAME normed input as attention (d.xn = norm(x0), which +// survives attention since recordOProj writes dx, not xn) and adds onto the residual in parallel — +// no second norm. Both leave dx += FFN(·). +func (d *Decoder) recordFFNSublayer(r recorder, b block, rows int) error { + if b.moeRouter != nil { + // sparse MoE — per BLOCK (DeepSeek-V2 mixes dense and MoE layers): pre-norm, then route → + // per-expert SwiGLU → weighted combine (+ optional shared expert). + if e := d.norm(r, d.dx.b, b.gFFN, b.bFFN, d.xn2.b, rows); e != nil { + return e + } + return d.recordMoE(r, b, d.xn2.b, rows) + } + if d.postNorm { + return d.recordFFN(r, b, d.dx.b, rows) // post-norm: FFN reads the raw residual; recordDownProj norms its output + } + if d.parallelRes { + return d.recordFFN(r, b, d.xn.b, rows) // parallel: FFN(norm(x0)), same norm as attn + } + if d.parallelTwoNorm { + return d.recordFFN(r, b, d.xn2.b, rows) // two-norm parallel: xn2 = norm2(x0), computed pre-attn-add + } + if e := d.norm(r, d.dx.b, b.gFFN, b.bFFN, d.xn2.b, rows); e != nil { + return e + } + if e := recordBarrier(r); e != nil { + return e + } + return d.recordFFN(r, b, d.xn2.b, rows) +} + +// recordMoE records the sparse Mixture-of-Experts FFN sublayer over the normed input `in`: +// gate logits = in·Router → top-k renormalized weights (MoEGate) → for every expert, SwiGLU(in) +// scaled by that expert's per-token weight and accumulated straight onto the residual dx (RowAxpy). +// This is DENSE eval (all experts run); the non-selected experts get weight 0, so it is exact — +// the sparse top-k gather is a later optimization. The batched command buffer stays valid because +// the routing weights are computed on-device, not by host control flow. +func (d *Decoder) recordMoE(r recorder, b block, in buffer, rows int) error { + e := b.moeRouter.record(r, in, d.moeGate.b, rows) // [rows, nExperts] router logits + e = firstErr(e, r.MoEGate(d.moeGate.b, d.moeW.b, rows, d.nExperts, d.topK, boolToInt(d.moeRaw), d.routedScale)) + sh, canFuse := r.(swiGLUHalvesRecorder) + for ex := range b.moeExperts { + xp := b.moeExperts[ex] + if xp.wGU != nil && canFuse { + // Column-fused expert gate|up: one gated GEMV into gu = [·, 2·hidden] then SwiGLU over its + // halves — halves the per-expert projection launches (the loop launches every expert, most + // gated-out but still launched). Gating zeros non-routed rows, so SwiGLUHalves(0)=0. Bit-exact. + e = firstErr(e, + xp.wGU.recordMoE(r, in, d.gu.b, rows, d.moeW.b, d.nExperts, ex), + sh.SwiGLUHalves(d.gu.b, d.gate.b, rows, d.hidden), + xp.wD.recordMoE(r, d.gate.b, d.mo.b, rows, d.moeW.b, d.nExperts, ex), + r.Copy2D(d.moeW.b, ex, d.nExperts, d.moeCol.b, 0, 1, rows, 1), + r.RowAxpy(d.dx.b, d.mo.b, d.moeCol.b, rows, d.d), + ) + continue + } + e = firstErr(e, + xp.wG.recordMoE(r, in, d.gate.b, rows, d.moeW.b, d.nExperts, ex), + xp.wU.recordMoE(r, in, d.up.b, rows, d.moeW.b, d.nExperts, ex), + d.binElem(r, d.gate.b, d.up.b, d.gate.b, binarySwiGLU, rows, d.hidden), // silu(gate)·up + xp.wD.recordMoE(r, d.gate.b, d.mo.b, rows, d.moeW.b, d.nExperts, ex), // expert output → mo + r.Copy2D(d.moeW.b, ex, d.nExperts, d.moeCol.b, 0, 1, rows, 1), // extract weight column [rows] + r.RowAxpy(d.dx.b, d.mo.b, d.moeCol.b, rows, d.d), // dx += w_ex · expert_ex + ) + } + if b.moeShared.wG != nil { + // Shared expert run on every token. Qwen2-MoE gates it by sigmoid(in·gate); DeepSeek-V2 runs + // it unconditionally (no gate), so its SwiGLU accumulates straight onto the residual. + s := b.moeShared + e = firstErr(e, + s.wG.record(r, in, d.gate.b, rows), + s.wU.record(r, in, d.up.b, rows), + d.binElem(r, d.gate.b, d.up.b, d.gate.b, binarySwiGLU, rows, d.hidden), // shared SwiGLU → gate + ) + if b.moeSharedGate != nil { // Qwen2-MoE: sigmoid-gated combine + e = firstErr(e, + s.wD.record(r, d.gate.b, d.mo.b, rows), + b.moeSharedGate.record(r, in, d.moeCol.b, rows), + r.Unary(d.moeCol.b, d.moeCol.b, unarySigmoid), + r.RowAxpy(d.dx.b, d.mo.b, d.moeCol.b, rows, d.d), // dx += sigmoid(gate)·shared_out + ) + } else { // DeepSeek-V2: ungated — dx += shared·Wdown directly + e = firstErr(e, s.wD.recordAdd(r, d.gate.b, d.mo.b, d.dx.b, rows, d.d)) + } + } + return e +} + +// recordMLAAttention records DeepSeek-V2 Multi-head Latent Attention for `rows` new tokens at +// position pos, leaving dx += attn·Wo. It fuses the reference's head-by-head build: input_layernorm → +// low-rank q/kv latent projections (each with its own RMSNorm) → decoupled RoPE on the per-head q_pe +// and the SHARED k_pe → assemble the rectangular per-head Q [qkHead] / K [qkHead] (k_pe broadcast to +// every head) and the VHead-wide V straight into the KV cache → MHARect (rectangular) → o_proj. +func (d *Decoder) recordMLAAttention(r recorder, b block, pos, rows int) error { + H, qkH, nope, rope, vH := d.h, d.qkHead, d.qkNope, d.qkRope, d.vHead + kvH := nope + vH // kv_b per-head width (k_nope + value) + qDim, vDim := H*qkH, H*vH + compW := d.kvLoRA + rope // kv_a output width: [kv_latent | k_pe] + + e := d.norm(r, d.dx.b, b.gAttn, b.bAttn, d.xn.b, rows) // input_layernorm → xn + + // Query: q = q_b_proj(q_a_layernorm(q_a_proj(xn))), then decoupled RoPE on each head's q_pe. + e = firstErr(e, + b.wqA.record(r, d.xn.b, d.mlaCQ.b, rows), + r.RMSNorm(d.mlaCQ.b, b.gQA, d.mlaCQ.b, rows, d.qLoRA, d.eps), + b.wqB.record(r, d.mlaCQ.b, d.mlaQ.b, rows), // [rows, H·qkH] + ) + for h := 0; h < H; h++ { + e = firstErr(e, r.RoPEAt(d.mlaQ.b, d.mlaInv.b, d.mlaQ.b, h*qkH+nope, rows, qDim, 1, rope, rope/2, pos, d.mlaPosDiv)) + } + + // KV: compressed = kv_a_proj(xn) → [kv_latent | k_pe]; RoPE the shared k_pe in place; norm the + // latent and expand kv = kv_b_proj(kv_a_layernorm(latent)). + e = firstErr(e, + b.wkvA.record(r, d.xn.b, d.mlaComp.b, rows), + r.RoPEAt(d.mlaComp.b, d.mlaInv.b, d.mlaComp.b, d.kvLoRA, rows, compW, 1, rope, rope/2, pos, d.mlaPosDiv), + r.Copy2D(d.mlaComp.b, 0, compW, d.mlaLatent.b, 0, d.kvLoRA, rows, d.kvLoRA), // extract kv_latent + r.RMSNorm(d.mlaLatent.b, b.gKvA, d.mlaLatent.b, rows, d.kvLoRA, d.eps), + b.wkvB.record(r, d.mlaLatent.b, d.mlaKV.b, rows), // [rows, H·kvH] + ) + + // Assemble the new tokens' K/V into the cache at row pos: K[h] = [k_nope_h | k_pe(shared)], + // V[h] = value_h. Copy2D writes rows-many rows at cache-row offset pos. + for h := 0; h < H; h++ { + e = firstErr(e, + r.Copy2D(d.mlaKV.b, h*kvH, H*kvH, b.kC, pos*qDim+h*qkH, qDim, rows, nope), // k_nope + r.Copy2D(d.mlaComp.b, d.kvLoRA, compW, b.kC, pos*qDim+h*qkH+nope, qDim, rows, rope), // shared k_pe → every head + r.Copy2D(d.mlaKV.b, h*kvH+nope, H*kvH, b.vC, pos*vDim+h*vH, vDim, rows, vH), // value + ) + } + + // Rectangular attention over the cache [0 : pos+rows], then the output projection onto dx. + return firstErr(e, + r.MHARect(d.mlaQ.b, b.kC, b.vC, d.mlaAttn.b, rows, pos+rows, H, H, qkH, vH, 1, 0, d.mlaScale), + b.wo.recordAdd(r, d.mlaAttn.b, d.mo.b, d.dx.b, rows, d.d), // dx += attn·Wo + ) +} + +// recordMambaBlock records ONE Mamba selective-state-space block for a single decode token +// (rows==1), leaving dx += mixer(RMSNorm(dx)). It mirrors nn.MambaBlock.Forward exactly, but the +// full-sequence OpConv1D/OpSSM are replaced by their per-timestep decode kernels (Conv1DStep, +// SSMStep) that carry the block's conv-window and SSM state across Step calls — no attention, no KV +// cache, O(1) work per token. The per-token recurrence is why Mamba records rows==1 only (StepN +// loops it): each token's state update depends on the previous token's. Only DtProj and the conv +// carry a bias (the HF checkpoint's other projections are bias-free); the Δ bias lands inside +// softplus. A = −exp(A_log) is precomputed on the host (b.mA). +func (d *Decoder) recordMambaBlock(r recorder, b block) error { + dI, N, K := d.dInner, d.mambaN, d.dConv + // Jamba inserts an RMSNorm on Δ_low/B/C between the x_proj split and the scan; plain Mamba passes + // nil gains, so jnorm is a no-op there (keeps this one recorder path serving both). + jnorm := func(buf, g buffer, w int) error { + if g == nil { + return nil + } + return r.RMSNorm(buf, g, buf, 1, w, d.eps) + } + // pre-norm → xn = RMSNorm(dx); the up-projection splits into the x branch and the z gate. + e := d.norm(r, d.dx.b, b.gAttn, nil, d.xn.b, 1) // RMSNorm (lnBias false for Mamba) + e = firstErr(e, + b.mInX.record(r, d.xn.b, d.mbXin.b, 1), // xin = xn·InX [dInner] + b.mInZ.record(r, d.xn.b, d.mbZ.b, 1), // z = xn·InZ [dInner] + // local mixing: one causal depthwise-conv step (state carried), then SiLU. + r.Conv1DStep(d.mbXin.b, b.mConvW, b.mConvB, b.mConvState, d.mbXc.b, dI, K), + r.Unary(d.mbXc.b, d.mbXc.b, unarySiLU), + // input-dependent Δ = softplus(dt_proj(RMSNorm?(dt_low_proj(xc))) + dt_bias). + b.mDtLow.record(r, d.mbXc.b, d.mbDtLow.b, 1), // Δ_low [dtRank] + jnorm(d.mbDtLow.b, b.mDtNorm, d.dtRank), // Jamba DtNorm + b.mDtProj.record(r, d.mbDtLow.b, d.mbDelta.b, 1), // → [dInner] + r.AddBias(d.mbDelta.b, b.mDtBias, d.mbDelta.b, 1, dI), + r.Unary(d.mbDelta.b, d.mbDelta.b, unarySoftplus), + // input-dependent SSM params B, C (Jamba RMSNorms each). + b.mBProj.record(r, d.mbXc.b, d.mbB.b, 1), // B [N] + jnorm(d.mbB.b, b.mBNorm, N), // Jamba BNorm + b.mCProj.record(r, d.mbXc.b, d.mbC.b, 1), // C [N] + jnorm(d.mbC.b, b.mCNorm, N), // Jamba CNorm + // selective scan: one recurrent timestep (state carried), y[dInner]. + r.SSMStep(d.mbXc.b, d.mbDelta.b, b.mA, d.mbB.b, d.mbC.b, b.mDskip, b.mSsmState, d.mbY.b, dI, N), + // gate y ⊙ SiLU(z), then down-project onto the residual: dx += y·OutProj. + r.Unary(d.mbZ.b, d.mbZ.b, unarySiLU), + r.Binary(d.mbY.b, d.mbZ.b, d.mbY.b, binaryMul), + b.mOutProj.recordAdd(r, d.mbY.b, d.mo.b, d.dx.b, 1, d.d), + ) + return e +} + +// resetMambaState zeroes every block's conv-window and SSM state — called at the start of a fresh +// sequence (pos==0), since the recurrence carries state across Step calls with no KV cache to reset. +func (d *Decoder) resetMambaState() error { + zc := make([]float32, d.dInner*(d.dConv-1)) + zs := make([]float32, d.dInner*d.mambaN) + for _, b := range d.blocks { + if b.mConvState == nil { // Jamba: attention layers have no conv/SSM state + continue + } + if e := b.mConvState.UploadF32(zc); e != nil { + return e + } + if e := b.mSsmState.UploadF32(zs); e != nil { + return e + } + } + return nil +} + +// recordJambaAttn records ONE Jamba attention mixer for a single decode token (rows==1), leaving +// dx += attn(RMSNorm(dx)). It is Llama-style grouped-query causal attention over the block's growing +// KV cache, but NoPE (no rotary): Jamba's Mamba layers carry position, so q/k are NEVER rotated — +// the only departure from the standard attention sublayer. Bias-free q/k/v/o. +func (d *Decoder) recordJambaAttn(r recorder, b block, pos int) error { + kvDim := d.kvDim + return firstErr( + d.recordAttnNorm(r, b, 1), // xn = RMSNorm(dx) (input_layernorm) + b.wq.record(r, d.xn.b, d.q.b, 1), // q/k/v = xn·Wq/Wk/Wv — NoPE, so no RoPE between here and the cache + b.wk.record(r, d.xn.b, d.k.b, 1), + b.wv.record(r, d.xn.b, d.v_.b, 1), + r.Blit(d.k.b, 0, b.kC, pos*kvDim, kvDim), // append this token's K/V rows to the cache + r.Blit(d.v_.b, 0, b.vC, pos*kvDim, kvDim), + d.recordMHA(r, d.q.b, b.kC, b.vC, d.attn.b, 1, pos+1), + d.recordOProj(r, b, 1), // dx += attn·Wo + ) +} + +// allocMambaScratch allocates the rows==1 Mamba decode scratch. Mamba has no attention or KV cache, +// so it skips the q/k/v/attn/qkv/gate/up buffers allocScratch would make; it keeps only the residual +// (dx), the pre-norm output (xn), the down-projection scratch (mo), the logits, and the per-block +// intermediates. Everything is one row wide — the sequential recurrence never batches. +func (d *Decoder) allocMambaScratch(mk func(data []float32) *bufSlot) { + dI, N, dtR := d.dInner, d.mambaN, d.dtRank + d.dx = mk(make([]float32, d.d)) + d.xn = mk(make([]float32, d.d)) + d.mo = mk(make([]float32, d.d)) + d.logits = mk(make([]float32, d.v)) + d.mbXin = mk(make([]float32, dI)) + d.mbZ = mk(make([]float32, dI)) + d.mbXc = mk(make([]float32, dI)) + d.mbDtLow = mk(make([]float32, dtR)) + d.mbDelta = mk(make([]float32, dI)) + d.mbB = mk(make([]float32, N)) + d.mbC = mk(make([]float32, N)) + d.mbY = mk(make([]float32, dI)) +} + +// recordMamba2Block records ONE Mamba-2 SSD block for a single decode token (rows==1), leaving +// dx += mixer(RMSNorm(dx)). It mirrors nlp.Mamba2Mixer.forward: fused in_proj → split [z | xBC | dt] +// → causal-conv step + SiLU over conv_dim → split [x | B | C] → Δ=softplus(dt+dt_bias) → the SSD scan +// (cu_ssd_step: scalar per-head decay a=exp(Δ·A), grouped B/C, +D·x skip) → gated RMSNorm norm(y·SiLU(z)) +// → out_proj onto the residual. A=−exp(A_log) precomputed on the host. Like Mamba-1 it is rows==1 +// (sequential recurrence). Splits are contiguous rows==1 sub-ranges, extracted with Blit. +func (d *Decoder) recordMamba2Block(r recorder, b block) error { + I, CD, H2 := d.m2Inter, d.m2CD, d.m2H + gN := d.m2G * d.m2N + // pre-norm → xn = RMSNorm(dx); fused in_proj → proj = [z(I) | xBC(CD) | dt(H2)]. + e := d.norm(r, d.dx.b, b.gAttn, nil, d.xn.b, 1) // RMSNorm (lnBias false) + e = firstErr(e, + b.m2InProj.record(r, d.xn.b, d.m2Proj.b, 1), + r.Blit(d.m2Proj.b, 0, d.m2Z.b, 0, I), // gate z + r.Blit(d.m2Proj.b, I, d.m2XBC.b, 0, CD), // conv input xBC + r.Blit(d.m2Proj.b, I+CD, d.m2Dt.b, 0, H2), // Δ pre-activation + // local mixing over conv_dim: causal depthwise-conv step (state carried) + SiLU. + r.Conv1DStep(d.m2XBC.b, b.m2ConvW, b.m2ConvB, b.m2ConvState, d.m2Xc.b, CD, d.m2Conv), + r.Unary(d.m2Xc.b, d.m2Xc.b, unarySiLU), + // split conv output → x(value, I) | B(gN) | C(gN). + r.Blit(d.m2Xc.b, 0, d.m2X.b, 0, I), + r.Blit(d.m2Xc.b, I, d.m2B.b, 0, gN), + r.Blit(d.m2Xc.b, I+gN, d.m2C.b, 0, gN), + // Δ = softplus(dt + dt_bias) [num_heads]. + r.AddBias(d.m2Dt.b, b.m2DtBias, d.m2Dt.b, 1, H2), + r.Unary(d.m2Dt.b, d.m2Dt.b, unarySoftplus), + // SSD scan: scalar per-head decay, grouped B/C, +D·x skip → y[I]. + r.SSDStep(d.m2X.b, d.m2Dt.b, b.m2A, d.m2B.b, d.m2C.b, b.m2Dskip, b.m2SsdState, d.m2Y.b, H2, d.m2P, d.m2G, d.m2N), + // gated RMSNorm: norm(y · SiLU(z)) · normW, over the intermediate width. + r.Unary(d.m2Z.b, d.m2Z.b, unarySiLU), + r.Binary(d.m2Y.b, d.m2Z.b, d.m2Y.b, binaryMul), + r.RMSNorm(d.m2Y.b, b.m2NormW, d.m2Y.b, 1, I, d.eps), + // out_proj onto the residual: dx += y · OutProj. + b.m2OutProj.recordAdd(r, d.m2Y.b, d.mo.b, d.dx.b, 1, d.d), + ) + return e +} + +// resetMamba2State zeroes every block's conv-window and SSD state (fresh sequence, pos==0). +func (d *Decoder) resetMamba2State() error { + zc := make([]float32, d.m2CD*(d.m2Conv-1)) + zs := make([]float32, d.m2H*d.m2N*d.m2P) + for _, b := range d.blocks { + if e := b.m2ConvState.UploadF32(zc); e != nil { + return e + } + if e := b.m2SsdState.UploadF32(zs); e != nil { + return e + } + } + return nil +} + +// allocMamba2Scratch allocates the rows==1 Mamba-2 decode scratch (no attention/KV buffers). +func (d *Decoder) allocMamba2Scratch(mk func(data []float32) *bufSlot) { + I, CD, H2 := d.m2Inter, d.m2CD, d.m2H + gN := d.m2G * d.m2N + projSize := I + CD + H2 + d.dx = mk(make([]float32, d.d)) + d.xn = mk(make([]float32, d.d)) + d.mo = mk(make([]float32, d.d)) + d.logits = mk(make([]float32, d.v)) + d.m2Proj = mk(make([]float32, projSize)) + d.m2Z = mk(make([]float32, I)) + d.m2XBC = mk(make([]float32, CD)) + d.m2Dt = mk(make([]float32, H2)) + d.m2Xc = mk(make([]float32, CD)) + d.m2X = mk(make([]float32, I)) + d.m2B = mk(make([]float32, gN)) + d.m2C = mk(make([]float32, gN)) + d.m2Y = mk(make([]float32, I)) +} + +// rwMix records the RWKV token-shift interpolation μ⊙x + (1−μ)⊙shift into out (rows==1). The +// reference forms it as shift + μ⊙(x−shift); with only add/mul recorder ops we use the algebraically +// identical μ⊙x + (1−μ)⊙shift, so mu and its precomputed complement omu are both uploaded. +func (d *Decoder) rwMixInto(r recorder, x, shift, mu, omu, out buffer) error { + return firstErr( + r.Binary(x, mu, d.rwT1.b, binaryMul), // μ⊙x + r.Binary(shift, omu, d.rwT2.b, binaryMul), // (1−μ)⊙shift + r.Binary(d.rwT1.b, d.rwT2.b, out, binaryAdd), + ) +} + +// recordRWKVBlock records ONE RWKV-4 block for a single decode token (rows==1), updating dx in place +// through both residuals. It mirrors nn.RWKVBlock.Step: a WKV TIME-MIX sublayer (LN1 → token-shift → +// r/k/v projections → the stabilized WKV recurrence via cu_wkv_step → receptance gate σ(r)⊙wkv → +// output projection → residual) and a gated squared-ReLU CHANNEL-MIX (LN2 → token-shift → σ(CWr·) +// gate over relu(CWk·)²·CWv → residual). No attention, no KV cache — the per-block token-shift and +// WKV states (rwPrevTM/rwPrevCM/rwAA/rwBB/rwPP) carry the whole history in O(1). Time-mix uses +// wq/wk/wv/wo as Wr/Wk/Wv/Wo; channel-mix uses wG/wD as CWk/CWv. +func (d *Decoder) recordRWKVBlock(r recorder, b block) error { + // TIME-MIX: xn = LN1(dx); r=σ(Wr·mix_r), k=Wk·mix_k, v=Wv·mix_v. + e := r.LayerNorm(d.dx.b, b.gAttn, b.bAttn, d.xn.b, 1, d.d, d.eps) + e = firstErr(e, + d.rwMixInto(r, d.xn.b, b.rwPrevTM, b.rwMuR, b.rwOmR, d.rwMix.b), + b.wq.record(r, d.rwMix.b, d.rwR.b, 1), + r.Unary(d.rwR.b, d.rwR.b, unarySigmoid), + d.rwMixInto(r, d.xn.b, b.rwPrevTM, b.rwMuK, b.rwOmK, d.rwMix.b), + b.wk.record(r, d.rwMix.b, d.rwK.b, 1), + d.rwMixInto(r, d.xn.b, b.rwPrevTM, b.rwMuV, b.rwOmV, d.rwMix.b), + b.wv.record(r, d.rwMix.b, d.rwV.b, 1), + // WKV recurrence (state carried), then receptance gate and output projection onto the residual. + r.WKVStep(d.rwK.b, d.rwV.b, b.rwW, b.rwU, b.rwAA, b.rwBB, b.rwPP, d.rwWkv.b, d.d), + r.Binary(d.rwWkv.b, d.rwR.b, d.rwWkv.b, binaryMul), // σ(r)⊙wkv (o==a: Binary clobbers b, not a) + b.wo.record(r, d.rwWkv.b, d.rwO.b, 1), + r.Binary(d.dx.b, d.rwO.b, d.dx.b, binaryAdd), // residual + r.Blit(d.xn.b, 0, b.rwPrevTM, 0, d.d), // prevTM ← xn (token-shift state) + ) + // CHANNEL-MIX: yn = LN2(dx); cr=σ(CWr·mix_r); ck=relu(CWk·mix_k)²; cv=CWv·ck; dx += cr⊙cv. + e = firstErr(e, + r.LayerNorm(d.dx.b, b.gFFN, b.bFFN, d.xn.b, 1, d.d, d.eps), + d.rwMixInto(r, d.xn.b, b.rwPrevCM, b.rwCMuR, b.rwOmCR, d.rwMix.b), + b.rwCWr.record(r, d.rwMix.b, d.rwCr.b, 1), + r.Unary(d.rwCr.b, d.rwCr.b, unarySigmoid), + d.rwMixInto(r, d.xn.b, b.rwPrevCM, b.rwCMuK, b.rwOmCK, d.rwMix.b), + b.wG.record(r, d.rwMix.b, d.rwCk.b, 1), // CWk: Dim→Hidden + r.Unary(d.rwCk.b, d.rwCk.b, unaryReLU2), // relu(·)² + b.wD.record(r, d.rwCk.b, d.rwCv.b, 1), // CWv: Hidden→Dim + r.Binary(d.rwCv.b, d.rwCr.b, d.rwCv.b, binaryMul), // σ(CWr)⊙cv (o==a: Binary clobbers b, not a) + r.Binary(d.dx.b, d.rwCv.b, d.dx.b, binaryAdd), // residual + r.Blit(d.xn.b, 0, b.rwPrevCM, 0, d.d), // prevCM ← yn + ) + return e +} + +// resetRWKVState zeroes every block's token-shift + WKV state for a fresh sequence (pos==0): prevTM/ +// prevCM/aa/bb start at 0 and pp at −1e38 (the pre-first-token WKV running-max, matching OpWKV). +func (d *Decoder) resetRWKVState() error { + zero := make([]float32, d.d) + ninf := make([]float32, d.d) + for i := range ninf { + ninf[i] = -1e38 + } + for _, b := range d.blocks { + for _, st := range []buffer{b.rwPrevTM, b.rwPrevCM, b.rwAA, b.rwBB} { + if e := st.UploadF32(zero); e != nil { + return e + } + } + if e := b.rwPP.UploadF32(ninf); e != nil { + return e + } + } + return nil +} + +// allocRWKVScratch allocates the rows==1 RWKV decode scratch (no attention/KV buffers). +func (d *Decoder) allocRWKVScratch(mk func(data []float32) *bufSlot) { + d.dx = mk(make([]float32, d.d)) + d.xn = mk(make([]float32, d.d)) + d.logits = mk(make([]float32, d.v)) + d.rwMix = mk(make([]float32, d.d)) + d.rwR = mk(make([]float32, d.d)) + d.rwK = mk(make([]float32, d.d)) + d.rwV = mk(make([]float32, d.d)) + d.rwWkv = mk(make([]float32, d.d)) + d.rwO = mk(make([]float32, d.d)) + d.rwT1 = mk(make([]float32, d.d)) + d.rwT2 = mk(make([]float32, d.d)) + d.rwCr = mk(make([]float32, d.d)) + d.rwCk = mk(make([]float32, d.rwHidden)) + d.rwCv = mk(make([]float32, d.d)) +} + +func (d *Decoder) recordFFN(r recorder, b block, in buffer, rows int) error { + if d.ffnGELU || d.ffnReLU2 { + act := unaryGELU // 2-layer MLP activation: GELU (GPT-NeoX/Phi/StarCoder2) or squared ReLU (Nemotron) + if d.ffnReLU2 { + act = unaryReLU2 + } + e := b.wG.record(r, in, d.gate.b, rows) // fc = in·Wfc + if b.fcBias != nil { + e = firstErr(e, r.AddBias(d.gate.b, b.fcBias, d.gate.b, rows, d.hidden)) + } + e = firstErr(e, + r.Unary(d.gate.b, d.gate.b, act), // act(fc) + d.recordDownProj(r, b, rows), // dx += act(fc)·Wproj + ) + if b.projBias != nil { + e = firstErr(e, r.AddBias(d.dx.b, b.projBias, d.dx.b, rows, d.d)) + } + return e + } + if d.ffnGEGLU { + // GeGLU (Gemma): down(GELU(gate)⊙up). Same 3-matrix gated shape as SwiGLU, but the gate + // activation is GELU and the elementwise product is a separate mul (no fused silu·b op). + if b.wGU != nil { + // Column-fused: one GEMV into gu = [·, 2·hidden], then GeGLU over its halves (GELU⊙up) into + // gate — one launch instead of two GEMVs + Unary GELU + Binary mul. Bit-identical. + e := b.wGU.record(r, in, d.gu.b, rows) + if gh, ok := r.(geGLUHalvesRecorder); ok { + e = firstErr(e, gh.GeGLUHalves(d.gu.b, d.gate.b, rows, d.hidden)) + } else { + e = firstErr(e, + r.Copy2D(d.gu.b, 0, 2*d.hidden, d.gate.b, 0, d.hidden, rows, d.hidden), + r.Copy2D(d.gu.b, d.hidden, 2*d.hidden, d.up.b, 0, d.hidden, rows, d.hidden), + r.Unary(d.gate.b, d.gate.b, unaryGELU), + d.binElem(r, d.gate.b, d.up.b, d.gate.b, binaryMul, rows, d.hidden), + ) + } + return firstErr(e, d.recordDownProj(r, b, rows)) + } + return firstErr( + b.wG.record(r, in, d.gate.b, rows), // gate = in·Wgate + b.wU.record(r, in, d.up.b, rows), // up = in·Wup + r.Unary(d.gate.b, d.gate.b, unaryGELU), // GELU(gate) + d.binElem(r, d.gate.b, d.up.b, d.gate.b, binaryMul, rows, d.hidden), // GELU(gate)⊙up + d.recordDownProj(r, b, rows), // dx += (·)·Wdown + ) + } + if b.wGU != nil { + // Column-fused ffn_gate|ffn_up: one GEMV into gu = [·, 2·hidden], then SwiGLU over its halves + // into gate — one launch instead of two GEMVs, bit-identical (each output column quantizes + // independently, so the fused weight equals the two separate ones concatenated). + e := b.wGU.record(r, in, d.gu.b, rows) + if sh, ok := r.(swiGLUHalvesRecorder); ok { + e = firstErr(e, sh.SwiGLUHalves(d.gu.b, d.gate.b, rows, d.hidden)) + } else { + // Backend without the fused op: split the gu halves out (strided) and SwiGLU as usual. + e = firstErr(e, + r.Copy2D(d.gu.b, 0, 2*d.hidden, d.gate.b, 0, d.hidden, rows, d.hidden), + r.Copy2D(d.gu.b, d.hidden, 2*d.hidden, d.up.b, 0, d.hidden, rows, d.hidden), + d.binElem(r, d.gate.b, d.up.b, d.gate.b, binarySwiGLU, rows, d.hidden), + ) + } + return firstErr(e, d.recordDownProj(r, b, rows)) // dx += swiglu·Wdown + } + e := b.wG.record(r, in, d.gate.b, rows) + e = firstErr(e, b.wU.record(r, in, d.up.b, rows)) + e = firstErr(e, recordBarrier(r)) + e = firstErr(e, d.binElem(r, d.gate.b, d.up.b, d.gate.b, binarySwiGLU, rows, d.hidden)) + e = firstErr(e, recordBarrier(r)) + return firstErr(e, d.recordDownProj(r, b, rows)) // dx += swiglu·Wdown +} + +// binElem records an elementwise op bounded to the rows*width elements actually in play. Decoder +// buffers are allocated for the whole context, so the unbounded form touches ctx*width elements for +// a single decoded row. Recorders without the capability fall back to the unbounded call, which is +// correct — the tail of the buffer holds stale values that no consumer reads — just far slower. +func (d *Decoder) binElem(r recorder, a, b, o buffer, op, rows, width int) error { + if bn, ok := r.(binaryNRecorder); ok && rows > 0 && width > 0 { + return bn.BinaryN(a, b, o, op, rows*width) + } + return r.Binary(a, b, o, op) +} + +// norm records the pre-sublayer normalization: LayerNorm-with-bias when lnBias (StableLM/Phi), +// else RMSNorm (Llama). gamma is the weight; beta is the LayerNorm bias (ignored under RMSNorm). +func (d *Decoder) norm(r recorder, x, gamma, beta, o buffer, rows int) error { + if d.lnBias { + return r.LayerNorm(x, gamma, beta, o, rows, d.d, d.eps) + } + return r.RMSNorm(x, gamma, o, rows, d.d, d.eps) +} + +// bFinalBeta is the final-norm LayerNorm bias, or nil under RMSNorm (bFinal unset). +func (d *Decoder) bFinalBeta() buffer { + if d.bFinal != nil { + return d.bFinal.b + } + return nil +} + +func (d *Decoder) ropeQK(r recorder, qkv, inv buffer, seq, stride, pos int) error { + if d.noRope { // MPT: position enters through the ALiBi attention bias, not rotary + return nil + } + if d.rotaryDim < d.dk { + return r.RoPEPartialPair(qkv, inv, seq, stride, d.h, 0, d.kvH, d.d, d.dk, d.rotaryDim, pos, d.posDiv) + } + return r.RoPEPair(qkv, inv, seq, stride, d.h, 0, d.kvH, d.d, d.dk, d.half, pos, d.posDiv) +} + +// allocScratch uploads the RoPE frequencies and allocates the per-step scratch buffers. They are +// sized for up to maxLen rows so StepN can process a whole prompt (or a speculative draft window) +// in one recorded step (§T418) — a few MB at typical configs. +func (d *Decoder) allocScratch(mk func(data []float32) *bufSlot) { + c := d.maxLen + d.dinv = mk(d.invHost) + d.dx = mk(make([]float32, c*d.d)) + d.xn = mk(make([]float32, c*d.d)) + d.xn2 = mk(make([]float32, c*d.d)) + d.q = mk(make([]float32, c*d.qDim)) + d.k = mk(make([]float32, c*d.kvDim)) + d.v_ = mk(make([]float32, c*d.kvDim)) + d.qkv = mk(make([]float32, c*(d.qDim+2*d.kvDim))) + d.attn = mk(make([]float32, c*d.qDim)) + d.ao = mk(make([]float32, c*d.d)) + d.gate = mk(make([]float32, c*d.hidden)) + d.up = mk(make([]float32, c*d.hidden)) + if d.ops.fusedGateUp { + d.gu = mk(make([]float32, c*2*d.hidden)) // fused gate|up GEMV output for the SwiGLUHalves path + } + d.mo = mk(make([]float32, c*d.d)) + d.logits = mk(make([]float32, c*d.v)) + if d.moe { + d.moeGate = mk(make([]float32, c*d.nExperts)) + d.moeW = mk(make([]float32, c*d.nExperts)) + d.moeCol = mk(make([]float32, c)) + } +} + +// mkBuf returns a bufSlot allocator that records the first error and tracks buffers for Release. +func (d *Decoder) mkBuf(err *error) func(data []float32) *bufSlot { + return func(data []float32) *bufSlot { + if *err != nil { + return &bufSlot{} + } + b, e := d.ops.newBuffer(data) + if e != nil { + *err = e + return &bufSlot{} + } + d.all = append(d.all, b) + return &bufSlot{b: b} + } +} + +// mkLin is the shared quantize-aware projection builder for the per-arch constructors: resident Q8_0 +// when the backend set ops.quantizeF32 (a NewXQ8CUDA entry point), else the f32 weight — byte-identical +// to the old inline `lin` closures when the hook is nil, so every existing f32 constructor is unchanged. +// One helper lets an arch gain Q8 decode by adding a hook-setting constructor, no per-arch code change. +func (d *Decoder) mkLin(mk func([]float32) *bufSlot, ops backendOps, errp *error) func(*tensor.Tensor) linear { + return func(w *tensor.Tensor) linear { + in, out := w.Shape()[0], w.Shape()[1] + if ops.quantizeF32 == nil { + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } + qw, qe := ops.quantizeF32(w) + if qe != nil { + if *errp == nil { + *errp = qe + } + return f32Linear{k: in, n: out} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } +} + +// f32Lin is the always-f32 projection builder, ignoring any Q8 hook. Used for weights whose +// quantization would be incorrect even under a NewXQ8CUDA path — notably the MoE router: Q8 rounding in +// the [dim,E] router logits could flip a top-k expert selection and change the output discontinuously, +// so routing stays exact while the (bandwidth-dominant) expert matrices go Q8. +func (d *Decoder) f32Lin(mk func([]float32) *bufSlot) func(*tensor.Tensor) linear { + return func(w *tensor.Tensor) linear { + in, out := w.Shape()[0], w.Shape()[1] + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } +} + +// mkLinS is mkLin for a weight pre-scaled by s (a per-element scalar folded into the upload — Granite's +// ResidualMult / logit-scale, Cohere's logit_scale). Q8_0 when ops.quantizeF32 is set (the scaled weight +// is quantized POST-scale), else f32 — byte-identical to the old inline f32-only linS when the hook is nil. +// Fixes a latent partial-Q8 bug: the per-arch f32-only linS left scaled lm_head / Wo / Wdown weights f32 +// even under a NewXQ8CUDA hook. +func (d *Decoder) mkLinS(mk func([]float32) *bufSlot, ops backendOps, errp *error) func(*tensor.Tensor, float32) linear { + return func(w *tensor.Tensor, s float32) linear { + in, out := w.Shape()[0], w.Shape()[1] + f := flat2D(w) + if s != 1 { + for i := range f { + f[i] *= s + } + } + if ops.quantizeF32 == nil { + return f32Linear{w: mk(f).b, k: in, n: out} + } + t := tensor.New(tensor.F32, tensor.Shape{in, out}) + copy(t.Storage().F32(), f) + qw, qe := ops.quantizeF32(t) + if qe != nil { + if *errp == nil { + *errp = qe + } + return f32Linear{k: in, n: out} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } +} + +// mkFused is mkLin for the fused QKV weight (§T613): the three [in,out] projections concatenated per +// input row into one [in, nq+nk+nv] matrix, then Q8_0-quantized whole (or f32). cu_qmatmul_q8 is a GEMM +// (rows>1), so batched-prefill StepN works; qkv bias / QK-norm / RoPE stay f32 and apply after. +func (d *Decoder) mkFused(mk func([]float32) *bufSlot, ops backendOps, errp *error) func(wq, wk, wv *tensor.Tensor) linear { + return func(wq, wk, wv *tensor.Tensor) linear { + in := wq.Shape()[0] + nq, nk, nv := wq.Shape()[1], wk.Shape()[1], wv.Shape()[1] + fq, fk, fv := flat2D(wq), flat2D(wk), flat2D(wv) + nt := nq + nk + nv + w := make([]float32, in*nt) + for i := range in { + row := w[i*nt : (i+1)*nt] + copy(row[:nq], fq[i*nq:(i+1)*nq]) + copy(row[nq:nq+nk], fk[i*nk:(i+1)*nk]) + copy(row[nq+nk:], fv[i*nv:(i+1)*nv]) + } + if ops.quantizeF32 == nil { + return f32Linear{w: mk(w).b, k: in, n: nt} + } + t := tensor.New(tensor.F32, tensor.Shape{in, nt}) + copy(t.Storage().F32(), w) + qw, qe := ops.quantizeF32(t) + if qe != nil { + if *errp == nil { + *errp = qe + } + return f32Linear{k: in, n: nt} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } +} + +// mkFused2 is mkFused for the fused gate|up weight: two [in,out] projections concatenated per input +// row into one [in, ng+nu] matrix, then quantized whole (or f32) — identical per-column quantization +// to the separate weights (each output column quantizes over its own K independently), so the fused +// GEMV output is bit-identical to two separate GEMVs. Used only when ops.fusedGateUp (CUDA SwiGLU). +func (d *Decoder) mkFused2(mk func([]float32) *bufSlot, ops backendOps, errp *error) func(wg, wu *tensor.Tensor) linear { + return func(wg, wu *tensor.Tensor) linear { + in := wg.Shape()[0] + ng, nu := wg.Shape()[1], wu.Shape()[1] + fg, fu := flat2D(wg), flat2D(wu) + nt := ng + nu + w := make([]float32, in*nt) + for i := range in { + row := w[i*nt : (i+1)*nt] + copy(row[:ng], fg[i*ng:(i+1)*ng]) + copy(row[ng:], fu[i*nu:(i+1)*nu]) + } + if ops.quantizeF32 == nil { + return f32Linear{w: mk(w).b, k: in, n: nt} + } + t := tensor.New(tensor.F32, tensor.Shape{in, nt}) + copy(t.Storage().F32(), w) + qw, qe := ops.quantizeF32(t) + if qe != nil { + if *errp == nil { + *errp = qe + } + return f32Linear{k: in, n: nt} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } +} + +// gateUpBuilder returns a helper that, per block, either column-fuses ffn_gate|ffn_up into one wGU +// weight (when ops.fusedGateUp — the CUDA SwiGLU/GeGLU path) or builds the separate wG/wU pair. Lets +// every SwiGLU/GeGLU builder opt into the fused FFN with one line, keeping non-CUDA backends identical. +func (d *Decoder) gateUpBuilder(mk func([]float32) *bufSlot, ops backendOps, errp *error) func(wg, wu *tensor.Tensor) (gu, g, u linear) { + lin := d.mkLin(mk, ops, errp) + fuseGU := d.mkFused2(mk, ops, errp) + return func(wg, wu *tensor.Tensor) (gu, g, u linear) { + if ops.fusedGateUp { + return fuseGU(wg, wu), nil, nil + } + return nil, lin(wg), lin(wu) + } +} + +// newDecoder uploads m's f32 weights via ops into device buffers and prepares a KV cache up to Ctx. +func newDecoder(m *nlp.Llama, ops backendOps) (*Decoder, error) { + cfg := m.Config + // A blockless model is an ERROR, never a legitimate degenerate case. The upload loop below is + // `for _, b := range m.Blocks`, so with no blocks it simply never runs and Step falls through + // embedding → final norm → lm_head: plausible logits, a full Generate sequence, err == nil + // everywhere, and the entire transformer stack silently skipped — the confident wrong answer + // §V29 forbids. Same guard as newDeepSeekV2Decoder / newMambaDecoder / newRWKVDecoder and the + // MoE builders; newDecoder (behind New, NewCUDA, NewVulkan, NewLlamaQ8CUDA, NewLlamaQ4KCUDA and + // the Qwen2/Qwen3/Phi3/Granite constructors) was the only sibling missing it. + if len(m.Blocks) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Llama has no blocks", ops.name) + } + // The same silence one step milder: a config claiming more (or fewer) layers than the checkpoint + // carries decodes the blocks it has and says nothing about the ones it does not. Nothing here + // reads cfg.Layers, so a config that leaves it unset (0) is not a claim and stays accepted — + // only an explicit disagreement is rejected. + if cfg.Layers > 0 && cfg.Layers != len(m.Blocks) { + return nil, fmt.Errorf("llamagpu(%s): config claims %d layers but model has %d blocks", ops.name, cfg.Layers, len(m.Blocks)) + } + d, derr := newDecoderCommon(cfg, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + + var err error + mk := d.mkBuf(&err) + // qOrF32 is the projection core: resident Q8_0 when the backend set ops.quantizeF32 (the + // NewLlamaQ8CUDA / NewQwen2Q8CUDA / NewQwen3Q8CUDA / NewPhi3Q8CUDA family), else the f32 weight. + // EVERY projection — fused QKV, o_proj, gate/up/down, lm_head — flows through it, so an f32 + // checkpoint gets a direct Q8 decode path (~4× less projection weight bandwidth) without a ggml + // conversion — and it reaches the Qwen2 (qkv-bias) / Qwen3 (QK-norm) variants that nlp.QuantLlama + // (Llama-only) cannot represent. Biases and QK-norm gains stay f32 and apply after the Q8 matmul. + qOrF32 := func(f []float32, in, out int) linear { + if ops.quantizeF32 == nil { + return f32Linear{w: mk(f).b, k: in, n: out} + } + t := tensor.New(tensor.F32, tensor.Shape{in, out}) + copy(t.Storage().F32(), f) + qw, qe := ops.quantizeF32(t) + if qe != nil { + if err == nil { + err = qe + } + return f32Linear{k: in, n: out} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } + lin := func(w *tensor.Tensor) linear { // Q8_0 (quantized build) or f32 [in,out] device weight + in, out := w.Shape()[0], w.Shape()[1] + return qOrF32(flat2D(w), in, out) + } + // IBM Granite config scalars (all identity when 0 or 1, so no-ops for Llama/Qwen/Mistral). Each + // folds into the upload rather than a runtime op: AttentionMult overrides the pre-softmax scale, + // EmbeddingMult scales gathered embeddings (d.embMult), ResidualMult scales each sublayer output + // (baked into Wo and Wdown, since their matmul IS the residual-add epilogue — Granite carries no + // projection bias, so nothing else in the add needs scaling), and LogitsScale divides the logits + // (baked as 1/LogitsScale into the untied lm_head). Separate uploads, so tied embed/lm_head is fine. + if cfg.AttentionMult != 0 { + d.scale = float32(cfg.AttentionMult) + } + if cfg.EmbeddingMult != 0 && cfg.EmbeddingMult != 1 { + d.embMult = float32(cfg.EmbeddingMult) + } + resMult := float32(1) + if cfg.ResidualMult != 0 && cfg.ResidualMult != 1 { + resMult = float32(cfg.ResidualMult) + } + outScale := float32(1) + if cfg.LogitsScale != 0 && cfg.LogitsScale != 1 { + outScale = float32(1 / cfg.LogitsScale) + } + linS := func(w *tensor.Tensor, s float32) linear { // weight with every element pre-scaled by s (then Q8 ∨ f32) + in, out := w.Shape()[0], w.Shape()[1] + f := flat2D(w) + if s != 1 { + for i := range f { + f[i] *= s + } + } + return qOrF32(f, in, out) + } + // fused QKV weight (§T613): weights are [in,out] with out along the row, so the fusion + // concatenates the three output bands PER INPUT ROW — one [in, D+2·kvDim] matrix whose + // product row is [q | k | v]. The separate wq/wk/wv uploads are dropped (no dup storage). + fused := func(wq, wk, wv *tensor.Tensor) linear { + in := wq.Shape()[0] + nq, nk, nv := wq.Shape()[1], wk.Shape()[1], wv.Shape()[1] + fq, fk, fv := flat2D(wq), flat2D(wk), flat2D(wv) + nt := nq + nk + nv + w := make([]float32, in*nt) + for i := range in { + row := w[i*nt : (i+1)*nt] + copy(row[:nq], fq[i*nq:(i+1)*nq]) + copy(row[nq:nq+nk], fk[i*nk:(i+1)*nk]) + copy(row[nq+nk:], fv[i*nv:(i+1)*nv]) + } + return qOrF32(w, in, nt) + } + // fused q/k/v bias [Bq | Bk | Bv] = [D+2·kvDim], for the Qwen2/Qwen2.5 family (q/k/v carry a + // projection bias, o_proj does not); nil for Llama/Mistral, which leaves recordQKVProj's + // AddBias unrecorded so those models stay byte-identical. + fusedBias := func(bq, bk, bv *tensor.Tensor) buffer { + if bq == nil && bk == nil && bv == nil { + return nil + } + fb := append(append(append([]float32{}, flat1D(bq)...), flat1D(bk)...), flat1D(bv)...) + return mk(fb).b + } + // per-head query/key RMSNorm gain [dk] (Qwen3); nil for Llama/Qwen2. Presence flips d.qkNorm, + // which arms recordQKNorm between the QKV projection and RoPE. + qkGain := func(n *nn.RMSNorm) buffer { + if n == nil { + return nil + } + d.qkNorm = true + return mk(flat1D(n.Gamma)).b + } + // gateUp column-fuses ffn_gate|ffn_up into one wGU weight when the backend can SwiGLU-over-halves + // (CUDA), else keeps the separate wG/wU pair. Same quantization path as the QKV fusion (qOrF32). + fuseGU := d.mkFused2(mk, ops, &err) + gateUp := func(wg, wu *tensor.Tensor) (gu, g, u linear) { + if ops.fusedGateUp { + return fuseGU(wg, wu), nil, nil + } + return nil, lin(wg), lin(wu) + } + //perfscan:ignore PS3032 model-build weight-upload loop over blocks, one-time + for _, b := range m.Blocks { + wgu, wg, wu := gateUp(b.FFN.Wgate, b.FFN.Wup) + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), qkvBias: fusedBias(b.Bq, b.Bk, b.Bv), wo: linS(b.Wo, resMult), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + qN: qkGain(b.QNorm), kN: qkGain(b.KNorm), + wGU: wgu, wG: wg, wU: wu, wD: linS(b.FFN.Wdown, resMult), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = linS(m.Out, outScale) // LogitsScale folded into the lm_head (1/LogitsScale) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newStableLMDecoder builds a device Decoder for an nlp.StableLM: the Llama-shaped SwiGLU/GQA +// core, but with LayerNorm-with-bias (lnBias) and PARTIAL rotary (rotaryDim < dk). StableLM has +// no attention bias and its rotary is split-half (matches OpRoPE), so weights upload directly — +// only the norms carry a bias and the rope is partial. Currently cuda-only (metal/vulkan +// RoPEPartialPair are unimplemented). §T757-class new-architecture GPU decode. +func newStableLMDecoder(m *nlp.StableLM, ops backendOps) (*Decoder, error) { + cfg := m.Config + // StableLMConfig's kvHeads()/headDim()/rotaryDim() are unexported, so replicate here. + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.HeadDim + if headDim <= 0 { + headDim = cfg.Dim / cfg.Heads + } + rotaryDim := cfg.RotaryDim + if rotaryDim <= 0 { + pct := cfg.RotaryPct + if pct == 0 { + pct = 0.25 + } + rotaryDim = int(float64(headDim) * pct) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + // LayerNorm-bias + partial rotary: rebuild the inverse-frequency table over rotaryDim + // (newDecoderCommon sized it to dk for full rope). + d.lnBias = true + d.rotaryDim = rotaryDim + if d.rotaryDim <= 0 || d.rotaryDim > d.dk || d.rotaryDim%2 != 0 { + return nil, fmt.Errorf("llamagpu(%s): StableLM rotaryDim %d invalid for headDim %d", ops.name, d.rotaryDim, d.dk) + } + invF64, posDiv64 := backend.RoPEFreqs(d.rotaryDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, d.rotaryDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + //perfscan:ignore PS3032 newStableLMDecoder block-upload, one-time model build + for _, b := range m.Blocks { + wgu, wg, wu := gateUp(b.FFN.Wgate, b.FFN.Wup) + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + gFFN: mk(flat1D(b.PostAttnNorm.Gamma)).b, bFFN: mk(flat1D(b.PostAttnNorm.Beta)).b, + wGU: wgu, wG: wg, wU: wu, wD: lin(b.FFN.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.bFinal = mk(flat1D(m.Norm.Beta)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newCohereDecoder builds a device Decoder for an nlp.Cohere (Command-R). Cohere composes existing +// generalizations: ONE-norm parallel residual (parallelRes — input_layernorm feeds both attention +// and FFN, summed onto x0, exactly like Phi), weight-only mean-centered LayerNorm (lnBias with a +// zero β), SwiGLU, GQA and FULL rope. Its interleaved (GPT-J) rotary is already baked into a q/k +// weight-row permutation by CohereFromHF, so standard split-half RoPE reproduces it. The only extra +// is logit_scale: the tied lm_head's logits are multiplied by a constant, folded into the Out +// upload. cuda-only (parallel-residual + LayerNorm reuse the Phi plumbing). +func newCohereDecoder(m *nlp.Cohere, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + if cfg.HeadDim != 0 && cfg.HeadDim != cfg.Dim/cfg.Heads { + return nil, fmt.Errorf("llamagpu(%s): Cohere decoupled head_dim %d (≠ dim/heads %d) not supported on the GPU decoder yet", ops.name, cfg.HeadDim, cfg.Dim/cfg.Heads) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // weight-only mean-centered LayerNorm (β is a zero vector) + d.parallelRes = true // one-norm parallel residual: input_layernorm feeds attn AND FFN + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + logitScale := float32(1) + if cfg.LogitScale != 0 && cfg.LogitScale != 1 { + logitScale = float32(cfg.LogitScale) + } + linS := d.mkLinS(mk, ops, &err) + //perfscan:ignore PS3032 newCohereDecoder block-upload, one-time model build + for _, b := range m.Blocks { + wgu, wg, wu := gateUp(b.FFN.Wgate, b.FFN.Wup) + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + // gFFN/bFFN unused: parallel one-norm reuses the attn norm (d.xn) for the FFN. + wGU: wgu, wG: wg, wU: wu, wD: lin(b.FFN.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.bFinal = mk(flat1D(m.Norm.Beta)) + d.out = linS(m.Out, logitScale) // logit_scale folded into the tied lm_head + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newNemotronDecoder builds a device Decoder for an nlp.Nemotron. Nemotron is a SEQUENTIAL +// two-norm decoder (input_layernorm feeds attention, post_attention_layernorm feeds the MLP, each +// on its own residual add) with three departures from Llama, all reusing existing plumbing: +// LayerNorm1P (an ordinary mean-centered LayerNorm — lnBias — whose γ=w+1 and β are folded in by +// NemotronFromHF), a squared-ReLU 2-layer MLP (ffnReLU2: down(relu²(up(x))), no gate, no bias), and +// PARTIAL rotary. cuda-only (partial rotary + the relu² unary are cuda-only). +func newNemotronDecoder(m *nlp.Nemotron, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.HeadDim + if headDim <= 0 { + headDim = cfg.Dim / cfg.Heads + } + if headDim != cfg.Dim/cfg.Heads { + return nil, fmt.Errorf("llamagpu(%s): Nemotron decoupled head_dim %d (≠ dim/heads %d) not supported on the GPU decoder yet", ops.name, headDim, cfg.Dim/cfg.Heads) + } + rotaryDim := cfg.RotaryDim + if rotaryDim <= 0 { + pct := cfg.RotaryPct + if pct == 0 { + pct = 0.5 + } + rotaryDim = int(float64(headDim) * pct) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // LayerNorm1P = mean-centered LayerNorm with a bias + d.ffnReLU2 = true // squared-ReLU 2-layer MLP + d.rotaryDim = rotaryDim + if d.rotaryDim <= 0 || d.rotaryDim > d.dk || d.rotaryDim%2 != 0 { + return nil, fmt.Errorf("llamagpu(%s): Nemotron rotaryDim %d invalid for headDim %d", ops.name, d.rotaryDim, d.dk) + } + invF64, posDiv64 := backend.RoPEFreqs(d.rotaryDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, d.rotaryDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + //perfscan:ignore PS3032 newNemotronDecoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + gFFN: mk(flat1D(b.PostAttnNorm.Gamma)).b, bFFN: mk(flat1D(b.PostAttnNorm.Beta)).b, + wG: lin(b.Wup), wD: lin(b.Wdown), // 2-layer relu² MLP: wG = up (fc), wD = down (proj); no gate + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.bFinal = mk(flat1D(m.Norm.Beta)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newGemmaDecoder builds a device Decoder for an nlp.Gemma (Gemma v1). Gemma is close to Llama — +// pre-norm RMSNorm, RoPE, GQA — with three departures, all reusing existing plumbing: the token +// embeddings are scaled by √dim right after lookup (d.embMult), RMSNorm uses (1+w) as the gain +// (folded in at load, so nn.RMSNorm applies directly), and the FFN is GeGLU rather than SwiGLU +// (ffnGEGLU — GELU gate + a plain elementwise product). The lm_head is tied to the embedding +// (Out = TokEmbᵀ). cuda-only for consistency with the rest of the new-arch family. +func newGemmaDecoder(m *nlp.Gemma, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.HeadDim + if headDim <= 0 { + headDim = cfg.Dim / cfg.Heads + } + // Gemma infers head_dim from q_proj (heads·head_dim rows) and it is DECOUPLED (head_dim need not + // equal dim/heads), so read it back from the loaded weights rather than trusting the config — + // GemmaFromHF may have left cfg.HeadDim at 0. + if qOut := m.Blocks[0].Wq.Shape()[1]; qOut%cfg.Heads == 0 { + headDim = qOut / cfg.Heads + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.FFN, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.ffnGEGLU = true // GELU-gated FFN + d.embMult = float32(math.Sqrt(float64(cfg.Dim))) // Gemma's √dim embedding normalizer + // Decoupled head_dim: override the geometry newDecoderCommon derived as dim/heads. qDim (the + // attention Q/O width) becomes h·head_dim, which for Gemma differs from the model dim d. + if headDim != d.dk { + d.dk = headDim + d.qDim = d.h * headDim + d.kvDim = d.kvH * headDim + d.half = headDim / 2 + d.rotaryDim = headDim + d.scale = float32(1.0 / math.Sqrt(float64(headDim))) + invF64, posDiv64 := backend.RoPEFreqs(headDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, headDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + } + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + // Column-fuse ffn_gate|ffn_up into wGU when the backend can GeGLU-over-halves (CUDA), else keep + // the separate wG/wU pair. GeGLU (Gemma) uses the same wGU layout as SwiGLU; only the halves op differs. + fuseGU := d.mkFused2(mk, ops, &err) + gateUp := func(wg, wu *tensor.Tensor) (gu, g, u linear) { + if ops.fusedGateUp { + return fuseGU(wg, wu), nil, nil + } + return nil, lin(wg), lin(wu) + } + //perfscan:ignore PS3032 newGemmaDecoder block-upload, one-time model build + for _, b := range m.Blocks { + wgu, wg, wu := gateUp(b.FFN.Wgate, b.FFN.Wup) + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + wGU: wgu, wG: wg, wU: wu, wD: lin(b.FFN.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + // Tied lm_head: Out = TokEmbᵀ. TokEmb is [vocab, dim]; the projection wants [dim, vocab]. + vocab, dim := m.TokEmb.Shape()[0], m.TokEmb.Shape()[1] + // Route the tied lm_head through the same quantize-aware builder so the Q8 constructors quantize it too + // — Gemma's vocab is huge (256k), so an f32 lm_head would dominate decode bandwidth and dilute Q8. The + // f32 path is byte-identical (mkLin with no hook == the old f32Linear{mk(outW)}). + outT := tensor.New(tensor.F32, tensor.Shape{dim, vocab}) + flat2DTInto(outT.Storage().F32()[:dim*vocab], m.TokEmb) + d.out = lin(outT) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newDeepSeekV2Decoder builds a device Decoder for an nlp.DeepSeekV2 — a stack of Multi-head Latent +// Attention blocks, each with a dense SwiGLU or a DeepSeekMoE sparse FFN (first_k_dense_replace +// split). MLA is the +// only new attention machinery (recordMLAAttention + rectangular MHARect); the FFN is a per-block +// dense SwiGLU or the DeepSeekMoE sparse FFN (raw-softmax·routedScale top-k routed experts + an +// ungated shared expert), dispatched by b.moeRouter. cuda-only (MHARect is cuda-only). +func newDeepSeekV2Decoder(m *nlp.DeepSeekV2, ops backendOps) (*Decoder, error) { + cfg := m.Config + if len(m.Blocks) == 0 { + return nil, fmt.Errorf("llamagpu(%s): DeepSeek-V2 has no blocks", ops.name) + } + qkHead := cfg.QKNope + cfg.QKRope + kvHead := cfg.QKNope + cfg.VHead + // hidden must cover the widest SwiGLU across every block (dense, routed-expert, shared) since the + // gate/up scratch is shared; nExperts/topK/routedScale come from the first MoE block if any. + hidden, nExperts, topK, hasMoE := 0, 0, 0, false + wider := func(sw *nn.SwiGLU) { + if sw != nil { + if h := sw.Wgate.Shape()[1]; h > hidden { + hidden = h + } + } + } + for _, b := range m.Blocks { + wider(b.Dense) + if b.MoE != nil { + hasMoE = true + for _, e := range b.MoE.Routed.Experts { + wider(e) + } + for _, s := range b.MoE.Shared { + wider(s) + } + if nExperts == 0 { + nExperts, topK = len(b.MoE.Routed.Experts), b.MoE.Routed.TopK + } + } + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: cfg.Heads, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.mla = true + if hasMoE { + d.moe, d.moeRaw = true, true // DeepSeek-V2 raw-softmax·routedScale gating + d.nExperts, d.topK = nExperts, topK + d.routedScale = 1 + if cfg.RoutedScale > 0 { + d.routedScale = float32(cfg.RoutedScale) + } + } + d.qkNope, d.qkRope, d.vHead, d.qkHead = cfg.QKNope, cfg.QKRope, cfg.VHead, qkHead + d.qLoRA, d.kvLoRA = cfg.QLoraRank, cfg.KVLoraRank + scale := cfg.SoftmaxScale + if scale <= 0 { + scale = 1.0 / math.Sqrt(float64(qkHead)) + } + d.mlaScale = float32(scale) + // decoupled RoPE frequencies over QKRope (single-head convention, Heads:1) + invF64, posDiv64 := backend.RoPEFreqs(cfg.QKRope, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: 1}) + d.mlaPosDiv = float32(posDiv64) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + // MLA scratch + the QKRope frequency table. + c := d.maxLen + mlaInv := make([]float32, cfg.QKRope/2) + for i := range mlaInv { + mlaInv[i] = float32(invF64[i]) + } + d.mlaInv = mk(mlaInv) + d.mlaCQ = mk(make([]float32, c*cfg.QLoraRank)) + d.mlaQ = mk(make([]float32, c*cfg.Heads*qkHead)) + d.mlaComp = mk(make([]float32, c*(cfg.KVLoraRank+cfg.QKRope))) + d.mlaLatent = mk(make([]float32, c*cfg.KVLoraRank)) + d.mlaKV = mk(make([]float32, c*cfg.Heads*kvHead)) + d.mlaAttn = mk(make([]float32, c*cfg.Heads*cfg.VHead)) + + //perfscan:ignore PS3032 newDeepSeekV2Decoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqA: lin(b.WqA), wqB: lin(b.WqB), wkvA: lin(b.WkvA), wkvB: lin(b.WkvB), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, gFFN: mk(flat1D(b.PostAttnNorm.Gamma)).b, + gQA: mk(flat1D(b.QANorm.Gamma)).b, gKvA: mk(flat1D(b.KvANorm.Gamma)).b, + // MLA cache: K is heads·qkHead wide, V is heads·VHead wide (rectangular). + kC: mk(make([]float32, c*cfg.Heads*qkHead)).b, vC: mk(make([]float32, c*cfg.Heads*cfg.VHead)).b, + } + if b.Dense != nil { // dense SwiGLU FFN + gb.wG, gb.wU, gb.wD = lin(b.Dense.Wgate), lin(b.Dense.Wup), lin(b.Dense.Wdown) + } else { // DeepSeekMoE: routed top-k experts + ungated shared expert(s) + experts := make([]moeFFN, len(b.MoE.Routed.Experts)) + for i, ex := range b.MoE.Routed.Experts { + wgu, wg, wu := gateUp(ex.Wgate, ex.Wup) + experts[i] = moeFFN{wGU: wgu, wG: wg, wU: wu, wD: lin(ex.Wdown)} + } + gb.moeRouter, gb.moeExperts = d.f32Lin(mk)(b.MoE.Routed.Router.W), experts // router stays f32 (routing is selection-sensitive) + if len(b.MoE.Shared) > 0 { // realized as a single fused SwiGLU; ungated (moeSharedGate nil) + s := b.MoE.Shared[0] + gb.moeShared = moeFFN{wG: lin(s.Wgate), wU: lin(s.Wup), wD: lin(s.Wdown)} + } + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + d.out = lin(m.LmHead) + d.allocScratch(mk) // standard scratch (dx/xn/xn2/gate/up/mo/logits reused; Llama-shaped q/qkv unused) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newMambaDecoder builds a device Decoder for an nlp.Mamba — the FIRST non-transformer GPU decoder. +// Each layer is a pre-RMSNorm + selective-state-space mixer (nn.MambaBlock): no attention, no KV +// cache, no FFN. Decode is a linear-time recurrence, so the block records the per-timestep conv/SSM +// decode kernels (Conv1DStep/SSMStep, cuda-only) carrying per-block state across Step calls. Only +// DtProj and the conv carry a bias (the HF checkpoint is otherwise bias-free); A = −exp(A_log) is +// precomputed on the host. The LM head is tied (logits = hidden · Embedᵀ). cuda-only. +func newMambaDecoder(m *nlp.Mamba, ops backendOps) (*Decoder, error) { + cfg := m.Config + if len(m.Layers) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Mamba has no layers", ops.name) + } + mix0 := m.Layers[0].Mixer + // Mamba has no attention: Heads=1 satisfies newDecoderCommon's geometry (all attention/KV scratch + // is skipped by allocMambaScratch). Hidden is unused (no FFN). Ctx is a pos bound only — the + // recurrence has no cache, so allocMambaScratch never scales with it; pick a generous limit. + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: 1 << 20, Dim: cfg.DModel, Heads: 1, KVHeads: 1, + Layers: cfg.Layers, Hidden: 1, Eps: cfg.Eps, + } + d, derr := newDecoderCommon(lc, m.Embed, ops) + if derr != nil { + return nil, derr + } + d.mamba = true + d.dInner, d.mambaN, d.dConv, d.dtRank = mix0.DInner, mix0.N, mix0.DConv, mix0.DtRank + + var err error + mk := d.mkBuf(&err) + // lin builds a projection: the f32 default, or (Q8 mode) a resident Q8_0 GEMV. Mamba decode is + // weight-bandwidth-bound (~230 MB of f32 weights streamed per token), so quantizing the projections + // to Q8_0 cuts that ~4× — the dominant decode lever. Only the matmul weights go Q8; the tiny + // per-channel buffers (A, dt_bias, conv, dskip, norms) stay f32 (not GEMVs, negligible bandwidth). + lin := func(w *tensor.Tensor) linear { + in, out := w.Shape()[0], w.Shape()[1] + if ops.quantizeF32 != nil { + qw, qe := ops.quantizeF32(w) + if qe != nil { + if err == nil { + err = qe + } + return f32Linear{k: in, n: out} // err set → constructor bails before use + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } + //perfscan:ignore PS3032 newMambaDecoder block-upload, one-time model build + for _, ly := range m.Layers { + mx := ly.Mixer + // A = −exp(A_log), precomputed on the host (ALog is [dInner, N], row-major — the SSMStep layout). + alog := flat2D(mx.ALog) + A := make([]float32, len(alog)) + for i, v := range alog { + A[i] = float32(-math.Exp(float64(v))) + } + gb := block{ + gAttn: mk(flat1D(ly.Norm.Gamma)).b, // pre-mixer RMSNorm γ (stored in gAttn; recordMambaBlock norms dx→xn) + mInX: lin(mx.InX.W), mInZ: lin(mx.InZ.W), + mDtLow: lin(mx.DtLow.W), mDtProj: lin(mx.DtProj.W), + mBProj: lin(mx.BProj.W), mCProj: lin(mx.CProj.W), + mOutProj: lin(mx.OutProj.W), + mConvW: mk(flat2D(mx.ConvW)).b, mConvB: mk(flat1D(mx.ConvB)).b, + mA: mk(A).b, mDskip: mk(flat1D(mx.Dskip)).b, mDtBias: mk(flat1D(mx.DtProj.B)).b, + mConvState: mk(make([]float32, d.dInner*(d.dConv-1))).b, + mSsmState: mk(make([]float32, d.dInner*d.mambaN)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Head) // tied head [d_model, vocab]: logits = hidden · Embedᵀ + d.allocMambaScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newJambaDecoder builds a device Decoder for an nlp.Jamba — the hybrid stack that interleaves Mamba +// selective-scan mixers with NoPE grouped-query attention, and per layer runs either a sparse-MoE or a +// dense SwiGLU FFN. Each layer is pre-norm(input_layernorm) → mixer → pre-norm(pre_ff_layernorm) → FFN, +// residual-added. Because some layers are Mamba (state carried across tokens), the whole model decodes +// sequentially like Mamba (rows==1 Step; StepN loops Step). Attention layers keep a growing KV cache; +// Mamba layers keep conv/SSM state. The MoE router uses raw (non-renormalized) softmax weights. +func newJambaDecoder(m *nlp.Jamba, ops backendOps) (*Decoder, error) { + cfg := m.Config + if len(m.Layers) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Jamba has no layers", ops.name) + } + // Infer the shared dims from whichever layers carry them: a Mamba layer fixes the SSM geometry, + // any FFN fixes the hidden width, a MoE layer fixes the expert count. + var mix0 *nn.MambaBlock + var hidden, nExperts int + for _, ly := range m.Layers { + if ly.Mamba != nil && mix0 == nil { + mix0 = ly.Mamba.Block + } + if ly.Dense != nil && hidden == 0 { + hidden = ly.Dense.Wgate.Shape()[1] + } + if ly.MoE != nil { + if hidden == 0 { + hidden = ly.MoE.Experts[0].Wgate.Shape()[1] + } + if nExperts == 0 { + nExperts = len(ly.MoE.Experts) + } + } + } + if mix0 == nil { + return nil, fmt.Errorf("llamagpu(%s): Jamba has no Mamba layer", ops.name) + } + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + topK := cfg.TopK + if topK <= 0 { + topK = 2 + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: 1 << 16, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.jamba = true + d.noRope = true // NoPE attention; also skips the (unused) RoPE-frequency upload in allocScratch + d.dInner, d.mambaN, d.dConv, d.dtRank = mix0.DInner, mix0.N, mix0.DConv, mix0.DtRank + if nExperts > 0 { + d.moe, d.moeRaw = true, true // Jamba routes raw softmax weights (no top-k renormalization) + d.nExperts, d.topK, d.routedScale = nExperts, topK, 1 + } + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) // Q8_0 when ops.quantizeF32 set (NewJambaQ8CUDA), else f32 + //perfscan:ignore PS3032 newJambaDecoder block-upload, one-time model build + for _, ly := range m.Layers { + gb := block{ + gAttn: mk(flat1D(ly.InputNorm.Gamma)).b, // input_layernorm (mixer pre-norm) + gFFN: mk(flat1D(ly.PreFFNorm.Gamma)).b, // pre_ff_layernorm (FFN pre-norm) + } + if ly.Mamba != nil { // Mamba mixer: the Mamba projections + Jamba's dt/b/c RMSNorms + carried state + jm := ly.Mamba + mx := jm.Block + alog := flat2D(mx.ALog) + A := make([]float32, len(alog)) + for i, v := range alog { + A[i] = float32(-math.Exp(float64(v))) + } + gb.jMamba = true + gb.mInX, gb.mInZ = lin(mx.InX.W), lin(mx.InZ.W) + gb.mDtLow, gb.mDtProj = lin(mx.DtLow.W), lin(mx.DtProj.W) + gb.mBProj, gb.mCProj = lin(mx.BProj.W), lin(mx.CProj.W) + gb.mOutProj = lin(mx.OutProj.W) + gb.mConvW, gb.mConvB = mk(flat2D(mx.ConvW)).b, mk(flat1D(mx.ConvB)).b + gb.mA, gb.mDskip, gb.mDtBias = mk(A).b, mk(flat1D(mx.Dskip)).b, mk(flat1D(mx.DtProj.B)).b + gb.mConvState = mk(make([]float32, d.dInner*(d.dConv-1))).b + gb.mSsmState = mk(make([]float32, d.dInner*d.mambaN)).b + gb.mDtNorm, gb.mBNorm, gb.mCNorm = mk(flat1D(jm.DtNorm.Gamma)).b, mk(flat1D(jm.BNorm.Gamma)).b, mk(flat1D(jm.CNorm.Gamma)).b + } else { // NoPE grouped-query attention: bias-free q/k/v/o + a growing KV cache + at := ly.Attn + gb.wq, gb.wk, gb.wv, gb.wo = lin(at.Wq), lin(at.Wk), lin(at.Wv), lin(at.Wo) + gb.kC = mk(make([]float32, d.maxLen*d.kvDim)).b + gb.vC = mk(make([]float32, d.maxLen*d.kvDim)).b + } + if ly.MoE != nil { // sparse MoE FFN (raw-softmax top-k, no shared expert) + experts := make([]moeFFN, len(ly.MoE.Experts)) + for i, ex := range ly.MoE.Experts { + experts[i] = moeFFN{wG: lin(ex.Wgate), wU: lin(ex.Wup), wD: lin(ex.Wdown)} + } + gb.moeRouter, gb.moeExperts = d.f32Lin(mk)(ly.MoE.Router), experts // router stays f32 (routing is selection-sensitive) + } else { // dense SwiGLU FFN + gb.wG, gb.wU, gb.wD = lin(ly.Dense.Wgate), lin(ly.Dense.Wup), lin(ly.Dense.Wdown) + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Out) // untied lm_head [dim, vocab] + d.allocScratch(mk) // attention + FFN + (d.moe) MoE scratch, all maxLen-sized + // Mamba-mixer intermediates (rows==1): the only scratch allocScratch doesn't cover. + dI, N, dtR := d.dInner, d.mambaN, d.dtRank + d.mbXin, d.mbZ, d.mbXc = mk(make([]float32, dI)), mk(make([]float32, dI)), mk(make([]float32, dI)) + d.mbDtLow, d.mbDelta = mk(make([]float32, dtR)), mk(make([]float32, dI)) + d.mbB, d.mbC, d.mbY = mk(make([]float32, N)), mk(make([]float32, N)), mk(make([]float32, dI)) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newMamba2Decoder builds a device Decoder for an nlp.Mamba2 — the state-space-duality sibling of +// Mamba. Each layer is a pre-RMSNorm + SSD mixer: fused in_proj → [z | xBC | dt] → causal-conv + SiLU +// → [x | B | C] → Δ=softplus(dt+dt_bias) → the SSD scan (scalar per-head decay, B/C shared across a +// group, +D·x skip; via cu_ssd_step) → gated RMSNorm norm(y·SiLU(z)) → out_proj. No attention, no KV +// cache — decode is a linear-time recurrence, so it reuses the Mamba rows==1 sequential Step/StepN +// plumbing. in_proj/out_proj are torch [out,in] (transposed on upload); A=−exp(A_log); tied LM head. +func newMamba2Decoder(m *nlp.Mamba2, ops backendOps) (*Decoder, error) { + cfg := m.Config + if len(m.Layers) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Mamba2 has no layers", ops.name) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: 1 << 20, Dim: cfg.DModel, Heads: 1, KVHeads: 1, + Layers: cfg.Layers, Hidden: 1, Eps: cfg.Eps, + } + d, derr := newDecoderCommon(lc, m.Embed, ops) + if derr != nil { + return nil, derr + } + d.mamba2 = true + d.m2H, d.m2P, d.m2G, d.m2N = cfg.NumHeads, cfg.HeadDim, cfg.NGroups, cfg.N + d.m2Conv, d.m2Inter, d.m2CD = cfg.DConv, cfg.Intermediate, cfg.Intermediate+2*cfg.NGroups*cfg.N + + var err error + mk := d.mkBuf(&err) + // linT uploads a torch [out,in] weight transposed to the [in,out] the recorder matmul expects. + linT := func(w *tensor.Tensor) linear { + out, in := w.Shape()[0], w.Shape()[1] + if ops.quantizeF32 != nil { // Q8 mode: transpose the [out,in] weight to [in,out], then quantize (§ weight-bandwidth lever) + wT, te := w.Transpose(0, 1) // zero-copy view; NewResidentBQ8 materializes it via Contiguous() + if te != nil { + if err == nil { + err = te + } + return f32Linear{k: in, n: out} + } + qw, qe := ops.quantizeF32(wT) + if qe != nil { + if err == nil { + err = qe + } + return f32Linear{k: in, n: out} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } + return f32Linear{w: mk(flat2DT(w)).b, k: in, n: out} + } + //perfscan:ignore PS3032 newMamba2Decoder block-upload, one-time model build + for _, ly := range m.Layers { + mx := ly.Mixer + aLog := flat1D(mx.ALog) + A := make([]float32, len(aLog)) + for i, v := range aLog { + A[i] = float32(-math.Exp(float64(v))) // A = −exp(A_log), per head + } + gb := block{ + gAttn: mk(flat1D(ly.Norm.Gamma)).b, // pre-mixer RMSNorm γ + m2InProj: linT(mx.InProj), m2OutProj: linT(mx.OutProj), + m2ConvW: mk(flat2D(mx.ConvW)).b, m2ConvB: mk(flat1D(mx.ConvB)).b, + m2A: mk(A).b, m2DtBias: mk(flat1D(mx.DtBias)).b, m2Dskip: mk(flat1D(mx.D)).b, + m2NormW: mk(flat1D(mx.NormW)).b, + m2ConvState: mk(make([]float32, d.m2CD*(d.m2Conv-1))).b, + m2SsdState: mk(make([]float32, d.m2H*d.m2N*d.m2P)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + // tied head: m.Head is already [d_model, vocab] = Embedᵀ, my [in=d_model, out=vocab] orientation. + if ops.quantizeF32 != nil { + if qw, qe := ops.quantizeF32(m.Head); qe != nil { + if err == nil { + err = qe + } + } else { + d.qweights = append(d.qweights, qw) + d.out = quantLinear{w: qw} + } + } else { + d.out = f32Linear{w: mk(flat2D(m.Head)).b, k: cfg.DModel, n: cfg.Vocab} + } + d.allocMamba2Scratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newRWKVDecoder builds a device Decoder for an nlp.RWKV — GoAI's third recurrent family (after the +// Mamba SSMs). Each layer is an nn.RWKVBlock: a WKV time-mixing sublayer (token-shift → r/k/v → the +// stabilized WKV recurrence via cu_wkv_step → receptance gate → output proj) and a gated squared-ReLU +// channel-mixing sublayer, each with its own LayerNorm + residual. No attention, no KV cache, no +// positional encoding — decode is O(1) recurrent, reusing the rows==1 sequential Step/StepN plumbing. +// An extra ln0 LayerNorm transforms the embedding before block 0; the head is UNTIED. cuda-only. +func newRWKVDecoder(m *nlp.RWKV, ops backendOps) (*Decoder, error) { + cfg := m.Config + if len(m.Blocks) == 0 { + return nil, fmt.Errorf("llamagpu(%s): RWKV has no blocks", ops.name) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: 1 << 20, Dim: cfg.Dim, Heads: 1, KVHeads: 1, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, + } + d, derr := newDecoderCommon(lc, m.Embed, ops) + if derr != nil { + return nil, derr + } + d.rwkv = true + d.lnBias = true // final norm (ln_out) is a full LayerNorm + d.rwHidden = cfg.Hidden + + var err error + mk := d.mkBuf(&err) + lin := func(w *tensor.Tensor) linear { // nn.Linear.W is already [in,out] (transposed at HF load) + in, out := w.Shape()[0], w.Shape()[1] + if ops.quantizeF32 != nil { // Q8 mode: same weight-bandwidth lever as Mamba (time-mix/channel-mix GEMVs) + qw, qe := ops.quantizeF32(w) + if qe != nil { + if err == nil { + err = qe + } + return f32Linear{k: in, n: out} + } + d.qweights = append(d.qweights, qw) + return quantLinear{w: qw} + } + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } + // mu uploads the interpolator and its 1−μ complement (token-shift compose needs both). + comp := func(mu *tensor.Tensor) buffer { + v := flat1D(mu) + o := make([]float32, len(v)) + for i, x := range v { + o[i] = 1 - x + } + return mk(o).b + } + expv := func(t *tensor.Tensor) buffer { // per-channel decay w = exp(WLog) + v := flat1D(t) + o := make([]float32, len(v)) + for i, x := range v { + o[i] = float32(math.Exp(float64(x))) + } + return mk(o).b + } + zeros := func() buffer { return mk(make([]float32, d.d)).b } + //perfscan:ignore PS3032 newRWKVDecoder block-upload, one-time model build + for _, bl := range m.Blocks { + gb := block{ + gAttn: mk(flat1D(bl.LN1.Gamma)).b, bAttn: mk(flat1D(bl.LN1.Beta)).b, // LN1 + gFFN: mk(flat1D(bl.LN2.Gamma)).b, bFFN: mk(flat1D(bl.LN2.Beta)).b, // LN2 + wq: lin(bl.Wr.W), wk: lin(bl.Wk.W), wv: lin(bl.Wv.W), wo: lin(bl.Wo.W), // time-mix Wr/Wk/Wv/Wo + wG: lin(bl.CWk.W), wD: lin(bl.CWv.W), rwCWr: lin(bl.CWr.W), // channel-mix CWk/CWv/CWr + rwMuR: mk(flat1D(bl.MuR)).b, rwMuK: mk(flat1D(bl.MuK)).b, rwMuV: mk(flat1D(bl.MuV)).b, + rwCMuR: mk(flat1D(bl.CMuR)).b, rwCMuK: mk(flat1D(bl.CMuK)).b, + rwOmR: comp(bl.MuR), rwOmK: comp(bl.MuK), rwOmV: comp(bl.MuV), + rwOmCR: comp(bl.CMuR), rwOmCK: comp(bl.CMuK), + rwW: expv(bl.WLog), rwU: mk(flat1D(bl.U)).b, + rwPrevTM: zeros(), rwPrevCM: zeros(), rwAA: zeros(), rwBB: zeros(), rwPP: zeros(), + } + d.blocks = append(d.blocks, gb) + } + d.rwPreLNg = mk(flat1D(m.PreLN.Gamma)) // ln0 (before block 0) + d.rwPreLNb = mk(flat1D(m.PreLN.Beta)) + d.gFinal = mk(flat1D(m.LNOut.Gamma)) // final LayerNorm ln_out + d.bFinal = mk(flat1D(m.LNOut.Beta)) + d.out = lin(m.Head) // untied head [dim, vocab] + d.allocRWKVScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newFalconDecoder builds a device Decoder for an nlp.Falcon (falcon-7b class: parallel_attn=True, +// multi_query=True, new_decoder_architecture=False). Every departure is an existing generalization: +// SINGLE-NORM parallel residual (parallelRes — one input_layernorm feeds attention AND the MLP, both +// summed onto x0, like Cohere), full LayerNorm WITH bias (lnBias), a 2-layer GELU MLP (ffnGELU), and +// multi-query attention (MQA = GQA with one KV head, kvH=1, from the fused query_key_value split). +// Full split-half rope, no linear biases, untied lm_head. cuda-only (parallel residual + LayerNorm). +func newFalconDecoder(m *nlp.Falcon, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = 1 // Falcon MQA + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // LayerNorm with bias + d.ffnGELU = true // 2-layer GELU MLP + d.parallelRes = true // one-norm parallel residual: input_layernorm feeds attn AND MLP + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + //perfscan:ignore PS3032 newFalconDecoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + // gFFN/bFFN unused: parallel one-norm reuses the attn norm (d.xn) for the MLP. + wG: lin(b.Wh), wD: lin(b.Wout), // 2-layer GELU MLP: wG = dense_h_to_4h, wD = dense_4h_to_h + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + d.bFinal = mk(flat1D(m.FinalNorm.Beta)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newOLMo2Decoder builds a device Decoder for an nlp.OLMo2 (Allen AI). OLMo2 shares Llama's SwiGLU/ +// GQA/RoPE core but reshuffles the residual structure in two ways, both new gated generalizations: +// - POST-NORM (postNorm): there is no input_layernorm; each sublayer reads the RAW residual and +// its OUTPUT is normed (post_attention_layernorm / post_feedforward_layernorm) before the add. +// - FULL-WIDTH QK-norm (qkNorm+qkNormFull): an RMSNorm over the ENTIRE q_proj / k_proj output +// (not per-head like Qwen3), applied before RoPE. +// +// The lm_head is untied. cuda-only. +func newOLMo2Decoder(m *nlp.OLMo2, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.postNorm = true // sublayer output is normed before the residual add + d.qkNorm = true // RMSNorm on the q/k projections before RoPE... + d.qkNormFull = true // ...over the full width, not per head + // decoupled head_dim (read back from q_proj; harmless when head_dim == dim/heads) + if qOut := m.Blocks[0].Wq.Shape()[1]; qOut%cfg.Heads == 0 && qOut/cfg.Heads != d.dk { + hd := qOut / cfg.Heads + d.dk, d.qDim, d.kvDim, d.half, d.rotaryDim = hd, d.h*hd, d.kvH*hd, hd/2, hd + d.scale = float32(1.0 / math.Sqrt(float64(hd))) + invF64, posDiv64 := backend.RoPEFreqs(hd, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, hd/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + } + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + //perfscan:ignore PS3032 newOLMo2Decoder block-upload, one-time model build + for _, b := range m.Blocks { + wgu, wg, wu := gateUp(b.FFN.Wgate, b.FFN.Wup) + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + // post-norm: gAttn = post_attention_layernorm (on the attn output), gFFN = post_feedforward_layernorm. + gAttn: mk(flat1D(b.PostAttnNorm.Gamma)).b, gFFN: mk(flat1D(b.PostFFNNorm.Gamma)).b, + qN: mk(flat1D(b.QNorm.Gamma)).b, kN: mk(flat1D(b.KNorm.Gamma)).b, // full-width [qDim]/[kvDim] + wGU: wgu, wG: wg, wU: wu, wD: lin(b.FFN.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newGemma2Decoder builds a device Decoder for an nlp.Gemma2. Gemma2 extends Gemma with three more +// departures, all reusing the generalizations added for its siblings: +// - SANDWICH norms (sandwich): each sublayer is pre-normed AND its output post-normed — four +// (1+w) RMSNorms per block (input/post_attention, pre_feedforward/post_feedforward). +// - attention-logit SOFT-CAP (attnCap → MHACap): scaled scores pass through cap·tanh(·/cap) +// before the mask+softmax. +// - query_pre_attn_scalar: the pre-softmax scale is 1/√scalar (not 1/√head_dim in general). +// +// Plus Gemma's √dim embedding normalizer, (1+w) RMSNorm, GeGLU, decoupled head_dim and tied lm_head. +// The final-logit soft-cap is a MONOTONIC map, so it never changes the greedy argmax and is omitted +// (greedy generation is identical with or without it). cuda-only (the soft-cap kernel is cuda-only). +func newGemma2Decoder(m *nlp.Gemma2, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.HeadDim + if headDim <= 0 { + headDim = cfg.Dim / cfg.Heads + } + if qOut := m.Blocks[0].Wq.Shape()[1]; qOut%cfg.Heads == 0 { + headDim = qOut / cfg.Heads + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.FFN, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.ffnGEGLU = true + d.sandwich = true + d.embMult = float32(math.Sqrt(float64(cfg.Dim))) + if cfg.AttnLogitCap > 0 { + d.attnCap = float32(cfg.AttnLogitCap) + } + // decoupled head_dim geometry override (as Gemma) + if headDim != d.dk { + d.dk, d.qDim, d.kvDim, d.half, d.rotaryDim = headDim, d.h*headDim, d.kvH*headDim, headDim/2, headDim + invF64, posDiv64 := backend.RoPEFreqs(headDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, headDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + } + // query_pre_attn_scalar sets the pre-softmax scale to 1/√scalar (0 → head_dim = standard). + qpa := cfg.QueryPreAttnScalar + if qpa == 0 { + qpa = float64(headDim) + } + d.scale = float32(1.0 / math.Sqrt(qpa)) + + var err error + mk := d.mkBuf(&err) + fused := d.mkFused(mk, ops, &err) + lin := d.mkLin(mk, ops, &err) + //perfscan:ignore PS3032 newGemma2Decoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, gAttn2: mk(flat1D(b.PostAttnNorm.Gamma)).b, + gFFN: mk(flat1D(b.PreFFNNorm.Gamma)).b, gFFN2: mk(flat1D(b.PostFFNNorm.Gamma)).b, + wG: lin(b.FFN.Wgate), wU: lin(b.FFN.Wup), wD: lin(b.FFN.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + // tied lm_head: Out = TokEmbᵀ [dim, vocab] + vocab, dim := m.TokEmb.Shape()[0], m.TokEmb.Shape()[1] + // Route the tied lm_head through the same quantize-aware builder so the Q8 constructors quantize it too + // — Gemma's vocab is huge (256k), so an f32 lm_head would dominate decode bandwidth and dilute Q8. The + // f32 path is byte-identical (mkLin with no hook == the old f32Linear{mk(outW)}). + outT := tensor.New(tensor.F32, tensor.Shape{dim, vocab}) + flat2DTInto(outT.Storage().F32()[:dim*vocab], m.TokEmb) + d.out = lin(outT) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newMPTDecoder builds a device Decoder for an nlp.MPT (MosaicML). MPT is the first ALiBi decoder: +// position enters ONLY through a per-head linear bias on the attention scores (no RoPE, no positional +// embedding), so newMPTDecoder sets noRope and uploads the backend.ALiBiSlopes so recordMHA routes +// attention through MHAALiBi. Otherwise it is a sequential-residual block with weight-only LayerNorm +// (lnBias, β=0), standard MHA (no GQA), a bias-free 2-layer GELU MLP (ffnGELU) and a tied lm_head. +// cuda-only (the ALiBi attention kernel is cuda-only). +func newMPTDecoder(m *nlp.MPT, ops backendOps) (*Decoder, error) { + cfg := m.Config + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: cfg.Heads, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: 10000, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // weight-only LayerNorm (β = 0) + d.ffnGELU = true // bias-free 2-layer GELU MLP + d.noRope = true // ALiBi, no rotary + + var err error + mk := d.mkBuf(&err) + slopes := make([]float32, d.h) + for i, s := range backend.ALiBiSlopes(d.h) { + slopes[i] = float32(s) + } + d.aliBiSlopes = mk(slopes).b + + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + //perfscan:ignore PS3032 newMPTDecoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.Norm1.Gamma)).b, bAttn: mk(flat1D(b.Norm1.Beta)).b, + gFFN: mk(flat1D(b.Norm2.Gamma)).b, bFFN: mk(flat1D(b.Norm2.Beta)).b, + wG: lin(b.Wup), wD: lin(b.Wdown), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + d.bFinal = mk(flat1D(m.FinalNorm.Beta)) + d.out = lin(m.Out) // tied wteᵀ, provided [dim, vocab] by the loader + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newMixtralDecoder builds a device Decoder for an nlp.Mixtral (sparse Mixture-of-Experts). The +// attention/norm stack is the plain Llama core (RMSNorm, fused QKV, full rope, GQA, and the optional +// per-head QK-norm shared with Qwen3-MoE); only the FFN sublayer becomes a sparse MoE: a router +// scores the experts, MoEGate picks the renormalized top-k weights, and every expert's SwiGLU is +// weighted-combined onto the residual (recordMoE). Dense expert eval — correct but not yet the sparse +// top-k gather. cuda-only (the MoE routing/combine kernels are cuda-only). +func newMixtralDecoder(m *nlp.Mixtral, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + if len(m.Blocks) == 0 || m.Blocks[0].MoE == nil || len(m.Blocks[0].MoE.Experts) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Mixtral has no MoE experts", ops.name) + } + // Derive the expert count / top-k from the loaded SparseMoE (MixtralFromHF infers Experts and + // may not write it back into Config). + nExperts := len(m.Blocks[0].MoE.Experts) + topK := m.Blocks[0].MoE.TopK + if topK <= 0 { + topK = 2 + } + hidden := cfg.Hidden + if hidden <= 0 { // inferred, may not be written back to Config — read it from an expert's Wgate [dim, ffn] + hidden = m.Blocks[0].MoE.Experts[0].Wgate.Shape()[1] + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.moe = true + d.nExperts = nExperts + d.topK = topK + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + qkGain := func(n *nn.RMSNorm) buffer { // optional per-head QK-norm (Qwen3-MoE); nil for Mixtral + if n == nil { + return nil + } + d.qkNorm = true + return mk(flat1D(n.Gamma)).b + } + //perfscan:ignore PS3032 newMixtralDecoder block-upload, one-time model build + for _, b := range m.Blocks { + experts := make([]moeFFN, len(b.MoE.Experts)) + for i, ex := range b.MoE.Experts { + wgu, wg, wu := gateUp(ex.Wgate, ex.Wup) + experts[i] = moeFFN{wGU: wgu, wG: wg, wU: wu, wD: lin(ex.Wdown)} + } + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + qN: qkGain(b.QNorm), kN: qkGain(b.KNorm), + moeRouter: d.f32Lin(mk)(b.MoE.Router.W), moeExperts: experts, // router stays f32 (routing is selection-sensitive) + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newQwen2MoEDecoder builds a device Decoder for an nlp.Qwen2MoE. Qwen2-MoE is a routed sparse MoE +// PLUS a shared expert: the FFN output is sparse_moe(x) + sigmoid(x·gate)·shared_expert(x), a single +// SwiGLU run on every token and scaled by a per-token sigmoid gate (recordMoE's shared-expert tail). +// Attention is the Llama core with Qwen2's q/k/v projection biases (o_proj bias-free). cuda-only. +func newQwen2MoEDecoder(m *nlp.Qwen2MoE, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + if len(m.Blocks) == 0 || m.Blocks[0].MoE == nil || len(m.Blocks[0].MoE.Experts) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Qwen2MoE has no MoE experts", ops.name) + } + b0 := m.Blocks[0] + nExperts := len(b0.MoE.Experts) + topK := b0.MoE.TopK + if topK <= 0 { + topK = 2 + } + hidden := b0.MoE.Experts[0].Wgate.Shape()[1] + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.moe = true + d.nExperts = nExperts + d.topK = topK + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + fusedBias := func(bq, bk, bv *tensor.Tensor) buffer { + if bq == nil && bk == nil && bv == nil { + return nil + } + return mk(append(append(append([]float32{}, flat1D(bq)...), flat1D(bk)...), flat1D(bv)...)).b + } + //perfscan:ignore PS3032 newQwen2MoEDecoder block-upload, one-time model build + for _, b := range m.Blocks { + experts := make([]moeFFN, len(b.MoE.Experts)) + for i, ex := range b.MoE.Experts { + wgu, wg, wu := gateUp(ex.Wgate, ex.Wup) + experts[i] = moeFFN{wGU: wgu, wG: wg, wU: wu, wD: lin(ex.Wdown)} + } + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), qkvBias: fusedBias(b.Bq, b.Bk, b.Bv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + moeRouter: d.f32Lin(mk)(b.MoE.Router.W), moeExperts: experts, // router stays f32 (routing is selection-sensitive) + moeShared: moeFFN{wG: lin(b.Shared.Wgate), wU: lin(b.Shared.Wup), wD: lin(b.Shared.Wdown)}, + moeSharedGate: lin(b.SharedGate), + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newGraniteMoEDecoder builds a device Decoder for an nlp.GraniteMoE — the MoE sibling of dense +// Granite: a plain Llama attention core + a sparse Mixture-of-Experts FFN, wrapped in the four +// Granite config scalars (all identity when 0/1, so this is newMixtralDecoder + the load-time folds +// from newDecoder's Granite handling). AttentionMult overrides the softmax scale; EmbeddingMult +// scales the gathered embedding; ResidualMult is baked into Wo AND every expert's Wdown (their +// matmul feeds the residual); LogitsScale is baked as 1/scale into the untied lm_head. cuda-only. +func newGraniteMoEDecoder(m *nlp.GraniteMoE, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + if len(m.Blocks) == 0 || m.Blocks[0].MoE == nil || len(m.Blocks[0].MoE.Experts) == 0 { + return nil, fmt.Errorf("llamagpu(%s): GraniteMoE has no MoE experts", ops.name) + } + b0 := m.Blocks[0] + nExperts := len(b0.MoE.Experts) + topK := b0.MoE.TopK + if topK <= 0 { + topK = 2 + } + hidden := b0.MoE.Experts[0].Wgate.Shape()[1] + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.moe = true + d.nExperts = nExperts + d.topK = topK + if cfg.AttentionMult != 0 { + d.scale = float32(cfg.AttentionMult) + } + if cfg.EmbeddingMult != 0 && cfg.EmbeddingMult != 1 { + d.embMult = float32(cfg.EmbeddingMult) + } + resMult := float32(1) + if cfg.ResidualMult != 0 && cfg.ResidualMult != 1 { + resMult = float32(cfg.ResidualMult) + } + outScale := float32(1) + if cfg.LogitsScale != 0 && cfg.LogitsScale != 1 { + outScale = float32(1 / cfg.LogitsScale) + } + + var err error + mk := d.mkBuf(&err) + gateUp := d.gateUpBuilder(mk, ops, &err) + linS := d.mkLinS(mk, ops, &err) // ResidualMult / logit-scale fold — quantize-aware (Q8 under the hook) + fused := d.mkFused(mk, ops, &err) + //perfscan:ignore PS3032 newGraniteMoEDecoder block-upload, one-time model build + for _, b := range m.Blocks { + experts := make([]moeFFN, len(b.MoE.Experts)) + for i, ex := range b.MoE.Experts { + wgu, wg, wu := gateUp(ex.Wgate, ex.Wup) + experts[i] = moeFFN{wGU: wgu, wG: wg, wU: wu, wD: linS(ex.Wdown, resMult)} + } + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: linS(b.Wo, resMult), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + moeRouter: d.f32Lin(mk)(b.MoE.Router.W), moeExperts: experts, // router stays f32 (routing is selection-sensitive) + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = linS(m.Out, outScale) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newOLMoEDecoder builds a device Decoder for an nlp.OLMoE (Allen AI sparse-MoE). OLMoE is the +// pre-norm sparse-MoE composition: the plain Llama attention core with FULL-WIDTH q/k RMSNorm +// (qkNorm+qkNormFull, one RMSNorm over the whole q/k projection — same as OLMo2) plus a sparse +// Mixture-of-Experts FFN (recordMoE) and an untied lm_head. cuda-only (MoE + full-width QK-norm). +func newOLMoEDecoder(m *nlp.OLMoE, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + if len(m.Blocks) == 0 || m.Blocks[0].MoE == nil || len(m.Blocks[0].MoE.Experts) == 0 { + return nil, fmt.Errorf("llamagpu(%s): OLMoE has no MoE experts", ops.name) + } + b0 := m.Blocks[0] + nExperts := len(b0.MoE.Experts) + topK := b0.MoE.TopK + if topK <= 0 { + topK = 2 + } + hidden := cfg.Hidden + if hidden <= 0 { + hidden = b0.MoE.Experts[0].Wgate.Shape()[1] + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.moe = true + d.nExperts = nExperts + d.topK = topK + d.qkNorm = true // full-width q/k RMSNorm before RoPE... + d.qkNormFull = true // ...over the whole projection (OLMo2/OLMoE) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + gateUp := d.gateUpBuilder(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + //perfscan:ignore PS3032 newOLMoEDecoder block-upload, one-time model build + for _, b := range m.Blocks { + experts := make([]moeFFN, len(b.MoE.Experts)) + for i, ex := range b.MoE.Experts { + wgu, wg, wu := gateUp(ex.Wgate, ex.Wup) + experts[i] = moeFFN{wGU: wgu, wG: wg, wU: wu, wD: lin(ex.Wdown)} + } + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), wo: lin(b.Wo), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + qN: mk(flat1D(b.QNorm.Gamma)).b, kN: mk(flat1D(b.KNorm.Gamma)).b, // full-width [qDim]/[kvDim] + moeRouter: d.f32Lin(mk)(b.MoE.Router.W), moeExperts: experts, // router stays f32 (routing is selection-sensitive) + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newStarCoder2Decoder builds a device Decoder for an nlp.StarCoder2: LayerNorm-with-bias + +// biased q/k/v/o projections + a biased 2-layer GELU MLP (c_fc → GELU → c_proj) + FULL rope + +// GQA + untied head. Every one of those is a core generalization (lnBias, biased projections, +// ffnGELU); rotary stays full (rotaryDim == dk). cuda-only. +func newStarCoder2Decoder(m *nlp.StarCoder2, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // LayerNorm-with-bias + d.ffnGELU = true // 2-layer GELU MLP + // rotary stays full (rotaryDim == dk, the newDecoderCommon default) — no inv override. + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + fusedBias := func(bq, bk, bv *tensor.Tensor) buffer { // [Bq | Bk | Bv] = [D+2·kvDim] + fb := append(append(append([]float32{}, flat1D(bq)...), flat1D(bk)...), flat1D(bv)...) + return mk(fb).b + } + //perfscan:ignore PS3032 newStarCoder2Decoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), qkvBias: fusedBias(b.Bq, b.Bk, b.Bv), + wo: lin(b.Wo), oBias: mk(flat1D(b.Bo)).b, + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + gFFN: mk(flat1D(b.PostAttnNorm.Gamma)).b, bFFN: mk(flat1D(b.PostAttnNorm.Beta)).b, + wG: lin(b.Wfc), fcBias: mk(flat1D(b.Bfc)).b, + wD: lin(b.Wproj), projBias: mk(flat1D(b.Bproj)).b, + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.bFinal = mk(flat1D(m.Norm.Beta)) + d.out = lin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newPhiDecoder builds a device Decoder for an nlp.Phi: the hardest new-arch shape so far — ONE +// input LayerNorm feeding BOTH attention and the MLP (one-norm parallel residual), biased q/k/v/dense +// projections, a biased 2-layer GELU MLP (fc1 → GELU → fc2), PARTIAL rotary, MHA, a biased untied +// lm_head and a final LayerNorm. Every one of those is a core generalization (parallelRes, biased +// projections, ffnGELU, lnBias, rotaryDim, outBias). cuda-only. +func newPhiDecoder(m *nlp.Phi, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.HeadDim + if headDim <= 0 { + headDim = cfg.Dim / cfg.Heads + } + rotaryDim := cfg.RotaryDim + if rotaryDim <= 0 { + pct := cfg.RotaryPct + if pct == 0 { + pct = 0.25 + } + rotaryDim = int(float64(headDim) * pct) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // LayerNorm-with-bias + d.ffnGELU = true // 2-layer GELU MLP + d.parallelRes = true // one-norm parallel residual (input_layernorm feeds attn AND mlp) + d.rotaryDim = rotaryDim + if d.rotaryDim <= 0 || d.rotaryDim > d.dk || d.rotaryDim%2 != 0 { + return nil, fmt.Errorf("llamagpu(%s): Phi rotaryDim %d invalid for headDim %d", ops.name, d.rotaryDim, d.dk) + } + invF64, posDiv64 := backend.RoPEFreqs(d.rotaryDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, d.rotaryDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + fusedBias := func(bq, bk, bv *tensor.Tensor) buffer { + fb := append(append(append([]float32{}, flat1D(bq)...), flat1D(bk)...), flat1D(bv)...) + return mk(fb).b + } + //perfscan:ignore PS3032 newPhiDecoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), qkvBias: fusedBias(b.Bq, b.Bk, b.Bv), + wo: lin(b.Wdense), oBias: mk(flat1D(b.Bdense)).b, + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + // gFFN/bFFN unused: parallel one-norm reuses the attn norm (d.xn) for the MLP. + wG: lin(b.Wfc1), fcBias: mk(flat1D(b.Bfc1)).b, + wD: lin(b.Wfc2), projBias: mk(flat1D(b.Bfc2)).b, + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + d.bFinal = mk(flat1D(m.FinalNorm.Beta)) + d.out = lin(m.Out) + d.outBias = mk(flat1D(m.OutBias)) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newGPTNeoXDecoder uploads an nlp.GPTNeoX onto the batched Decoder core. GPT-NeoX is the TWO-norm +// parallel residual: a separate input_layernorm feeds attention and post_attention_layernorm feeds +// the MLP, but BOTH read the raw residual x0 and both outputs sum onto it. Beyond that it is +// LayerNorm-with-bias, biased q/k/v/dense, a biased GELU MLP, (partial or full) rotary and an +// untied lm_head with no bias. The fourth new-arch GPU graph decoder — and the one that exercises +// parallelTwoNorm (the FFN norm must be taken from x0 before the attention output is added back). +func newGPTNeoXDecoder(m *nlp.GPTNeoX, ops backendOps) (*Decoder, error) { + cfg := m.Config + kvH := cfg.KVHeads + if kvH <= 0 { + kvH = cfg.Heads + } + headDim := cfg.Dim / cfg.Heads + rotaryDim := cfg.RotaryDim + if rotaryDim <= 0 { + pct := cfg.RotaryPct + if pct == 0 { + pct = 1 // GPT-NeoX default is full rope + } + rotaryDim = int(float64(headDim) * pct) + } + lc := nlp.LlamaConfig{ + Vocab: cfg.Vocab, Ctx: cfg.Ctx, Dim: cfg.Dim, Heads: cfg.Heads, KVHeads: kvH, + Layers: cfg.Layers, Hidden: cfg.Hidden, Eps: cfg.Eps, RopeBase: cfg.RopeBase, + } + d, derr := newDecoderCommon(lc, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + d.lnBias = true // LayerNorm-with-bias + d.ffnGELU = true // 2-layer GELU MLP + d.parallelTwoNorm = true // two-norm parallel residual (input_ln → attn, post_attn_ln → mlp, both over x0) + d.rotaryDim = rotaryDim + if d.rotaryDim <= 0 || d.rotaryDim > d.dk || d.rotaryDim%2 != 0 { + return nil, fmt.Errorf("llamagpu(%s): GPT-NeoX rotaryDim %d invalid for headDim %d", ops.name, d.rotaryDim, d.dk) + } + invF64, posDiv64 := backend.RoPEFreqs(d.rotaryDim, backend.RoPEAttrs{Base: cfg.RopeBase, Heads: d.h}) + d.invHost = make([]float32, d.rotaryDim/2) + for i := range d.invHost { + d.invHost[i] = float32(invF64[i]) + } + d.posDiv = float32(posDiv64) + + var err error + mk := d.mkBuf(&err) + lin := d.mkLin(mk, ops, &err) + fused := d.mkFused(mk, ops, &err) + fusedBias := func(bq, bk, bv *tensor.Tensor) buffer { + fb := append(append(append([]float32{}, flat1D(bq)...), flat1D(bk)...), flat1D(bv)...) + return mk(fb).b + } + //perfscan:ignore PS3032 newGPTNeoXDecoder block-upload, one-time model build + for _, b := range m.Blocks { + gb := block{ + wqkv: fused(b.Wq, b.Wk, b.Wv), qkvBias: fusedBias(b.Bq, b.Bk, b.Bv), + wo: lin(b.Wo), oBias: mk(flat1D(b.Bo)).b, + gAttn: mk(flat1D(b.InputNorm.Gamma)).b, bAttn: mk(flat1D(b.InputNorm.Beta)).b, + gFFN: mk(flat1D(b.PostAttnNorm.Gamma)).b, bFFN: mk(flat1D(b.PostAttnNorm.Beta)).b, + wG: lin(b.Wh), fcBias: mk(flat1D(b.Bh)).b, + wD: lin(b.Wout), projBias: mk(flat1D(b.Bout)).b, + kC: mk(make([]float32, d.maxLen*d.kvDim)).b, vC: mk(make([]float32, d.maxLen*d.kvDim)).b, + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.FinalNorm.Gamma)) + d.bFinal = mk(flat1D(m.FinalNorm.Beta)) + d.out = lin(m.Out) // untied embed_out, no bias + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// newQuantDecoder uploads a quantized Llama: RMSNorm gains + device-resident KV caches, every +// projection as a RESIDENT quantized weight consumed by the record-mode QMatMulResident (§T415). +// Cache storage stays f32 unless the constructor supplies newF16KVBuffer; that opt-in is deliberately +// limited to dk=64, the Metal kernel shape proven by P-01M093C3MFF0E8E1THCWYA1757. +func newQuantDecoder(m *nlp.QuantLlama, ops backendOps) (*Decoder, error) { + if ops.uploadQWeight == nil { + return nil, fmt.Errorf("llamagpu(%s): backend has no resident quantized upload", ops.name) + } + cfg := m.Config + d, derr := newDecoderCommon(cfg, m.TokEmb, ops) + if derr != nil { + return nil, derr + } + if ops.newF16KVBuffer != nil { + if d.dk != 64 { + return nil, fmt.Errorf("llamagpu(%s): f16 KV cache requires head dimension 64, got %d", ops.name, d.dk) + } + d.f16KV = true + } + var err error + mk := d.mkBuf(&err) + mkCache := func(n int) buffer { + if !d.f16KV { + return mk(make([]float32, n)).b + } + if err != nil { + return nil + } + b, e := ops.newF16KVBuffer(n) + if e != nil { + err = e + return nil + } + d.all = append(d.all, b) + return b + } + qlin := func(q *nn.QuantLinear) linear { // resident quantized [Out,In] weight + if err != nil { + return quantLinear{} + } + w, e := ops.uploadQWeight(q.Weight, uint32(q.QT), q.Out, q.In) + if e != nil { + err = e + return quantLinear{} + } + d.qweights = append(d.qweights, w) + return quantLinear{w: w} + } + // fused QKV for quantized blocks (§T613): ggml resident weights are [Out,In] ROW-major, so + // fusing = appending the raw quantized bytes (out rows q‖k‖v) — valid only when all three + // projections share one quant type; mixed-type blocks keep the unfused three-matmul chain. + qfused := func(q1, q2, q3 *nn.QuantLinear) linear { + if err != nil || q1.QT != q2.QT || q2.QT != q3.QT { + return nil + } + wb := make([]byte, 0, len(q1.Weight)+len(q2.Weight)+len(q3.Weight)) + wb = append(append(append(wb, q1.Weight...), q2.Weight...), q3.Weight...) + w, e := ops.uploadQWeight(wb, uint32(q1.QT), q1.Out+q2.Out+q3.Out, q1.In) + if e != nil { + err = e + return nil + } + d.qweights = append(d.qweights, w) + return quantLinear{w: w} + } + qgroup := func(weights ...linear) linear { + if err != nil || ops.groupQWeights == nil || len(weights) != 3 { + return nil + } + qw := make([]qweight, 3) + for i, w := range weights { + ql, ok := w.(quantLinear) + if !ok || ql.w == nil { + return nil + } + qw[i] = ql.w + } + g, e := ops.groupQWeights(qw) + if e == backend.ErrQuantUnsupported { + return nil + } + if e != nil { + err = e + return nil + } + d.qweights = append(d.qweights, g) + return quantLinear{w: g} + } + // Column-fused ffn_gate|ffn_up for quantized blocks (same trick as qfused: ggml resident weights are + // [Out,In] row-major, so fusing = appending the raw quant bytes of gate‖up out-rows). Only when the + // backend supports SwiGLUHalves (ops.fusedGateUp) and both share one quant type; else nil → the + // unfused two-GEMV wG/wU path. Mirrors newDecoder's wGU so the raw-quant decode matches its speed. + qfused2 := func(q1, q2 *nn.QuantLinear) linear { + if err != nil || !ops.fusedGateUp || q1.QT != q2.QT { + return nil + } + wb := make([]byte, 0, len(q1.Weight)+len(q2.Weight)) + wb = append(append(wb, q1.Weight...), q2.Weight...) + w, e := ops.uploadQWeight(wb, uint32(q1.QT), q1.Out+q2.Out, q1.In) + if e != nil { + err = e + return nil + } + d.qweights = append(d.qweights, w) + return quantLinear{w: w} + } + //perfscan:ignore PS3032 newQuantDecoder block-upload, one-time model build + for _, b := range m.Blocks { + wgu := qfused2(b.FFN.Gate, b.FFN.Up) + gb := block{ + wqkv: qfused(b.Wq, b.Wk, b.Wv), wo: qlin(b.Wo), + gAttn: mk(flat1D(b.AttnNorm.Gamma)).b, gFFN: mk(flat1D(b.FFNNorm.Gamma)).b, + wGU: wgu, wD: qlin(b.FFN.Down), + kC: mkCache(d.maxLen * d.kvDim), vC: mkCache(d.maxLen * d.kvDim), + } + if wgu == nil { // no SwiGLUHalves backend or mixed types: keep the two-GEMV gate/up path + gb.wG, gb.wU = qlin(b.FFN.Gate), qlin(b.FFN.Up) + } + if gb.wqkv == nil { // mixed quant types: keep exact raw projections for single-token decode + gb.wq, gb.wk, gb.wv = qlin(b.Wq), qlin(b.Wk), qlin(b.Wv) + gb.wqkvPrefill = qgroup(gb.wq, gb.wk, gb.wv) + } + d.blocks = append(d.blocks, gb) + } + d.gFinal = mk(flat1D(m.Norm.Gamma)) + d.out = qlin(m.Out) + d.allocScratch(mk) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// encodeStep records the whole single-token layer stack + vocab head for absolute position +// pos into a fresh command buffer WITHOUT committing it. The recorded commands depend only +// on pos — not on the token, whose embedding is read from d.dx at execution time — which is +// what makes the §T614 encode-overlap legal: while the GPU executes step pos, the host +// pre-encodes step pos+1, and the next Step only uploads dx and commits. +func (d *Decoder) encodeStep(pos int) (recorder, error) { + newRecorder := d.ops.newRecorder + if d.ops.newDecodeRecorder != nil { + newRecorder = d.ops.newDecodeRecorder + } + r, err := newRecorder() + if err != nil { + return nil, err + } + D, H, KVH, dk, kvDim := d.qDim, d.h, d.kvH, d.dk, d.kvDim // D = attention Q/O width (h·dk); = d.d unless head_dim decoupled + if d.rwkv { // RWKV: an extra LayerNorm (ln0) transforms the embedding once before block 0 + if e := r.LayerNorm(d.dx.b, d.rwPreLNg.b, d.rwPreLNb.b, d.dx.b, 1, d.d, d.eps); e != nil { + r.Free() + return nil, e + } + } + for _, b := range d.blocks { + if d.rwkv { // RWKV-4 block: WKV time-mix + gated squared-ReLU channel-mix, no attention/KV cache + if e := d.recordRWKVBlock(r, b); e != nil { + r.Free() + return nil, e + } + continue + } + if d.mamba { // Mamba selective-scan block: no attention/FFN split — the mixer is the whole layer + if e := d.recordMambaBlock(r, b); e != nil { + r.Free() + return nil, e + } + continue + } + if d.mamba2 { // Mamba-2 SSD block (scalar-decay, grouped B/C, gated RMSNorm) + if e := d.recordMamba2Block(r, b); e != nil { + r.Free() + return nil, e + } + continue + } + if d.jamba { // hybrid: per-block mixer (Mamba OR NoPE attention), then a per-block MoE/dense FFN + var e error + if b.jMamba { + e = d.recordMambaBlock(r, b) // Mamba mixer w/ Jamba dt/b/c norms → dx += mixer(inputNorm(dx)) + } else { + e = d.recordJambaAttn(r, b, pos) // NoPE GQA attention → dx += attn(inputNorm(dx)) + } + if e = firstErr(e, d.recordFFNSublayer(r, b, 1)); e != nil { // MoE (raw softmax) or dense SwiGLU + r.Free() + return nil, e + } + continue + } + if d.mla { // DeepSeek-V2 Multi-head Latent Attention (rectangular, low-rank latent KV) + if e := firstErr(d.recordMLAAttention(r, b, pos, 1), d.recordFFNSublayer(r, b, 1)); e != nil { + r.Free() + return nil, e + } + continue + } + // attention projections: fused single QKV matmul when available (§T613 — the q/k + // bands rotate IN PLACE inside the combined row, k/v append to the cache straight + // from their bands, attention reads q at band offset 0), else the unfused three. + // Recording order is execution order: norm FIRST, then the projection branch. + e := d.recordAttnNorm(r, b, 1) + e = firstErr(e, recordBarrier(r)) + qBuf := d.q.b + if e == nil && b.wqkv != nil { + qBuf = d.qkv.b + e = firstErr( + d.recordQKVProj(r, b, 1), + d.recordQKNorm(r, b, 1, D+2*kvDim), // per-head QK-norm (Qwen3); no-op otherwise + ) + e = firstErr(e, recordBarrier(r)) + if fr, ok := r.(ropePairF16KVAppendRecorder); e == nil && ok && d.f16KV && + !d.noRope && d.rotaryDim == dk && dk == 64 { + e = fr.RoPEPairF16KVAppend( + d.qkv.b, d.dinv.b, b.kC, b.vC, D+2*kvDim, + H, 0, KVH, D, dk, d.half, D+kvDim, kvDim, pos, pos*kvDim, d.posDiv, + ) + } else if e == nil { + e = firstErr( + d.ropeQK(r, d.qkv.b, d.dinv.b, 1, D+2*kvDim, pos), // q+k bands, ONE dispatch (§T613); full ∨ partial rotary + d.recordKVPairBlit(r, d.qkv.b, D, d.qkv.b, D+kvDim, b.kC, b.vC, pos*kvDim, kvDim), + ) + } + } else if e == nil { + e = firstErr( + b.wq.record(r, d.xn.b, d.q.b, 1), + b.wk.record(r, d.xn.b, d.k.b, 1), + b.wv.record(r, d.xn.b, d.v_.b, 1), + ) + e = firstErr(e, recordBarrier(r)) + if fr, ok := r.(ropeF16KVAppendRecorder); e == nil && ok && d.f16KV && + !d.noRope && d.rotaryDim == dk && dk == 64 { + e = fr.RoPEF16KVAppend( + d.q.b, d.k.b, d.v_.b, d.dinv.b, b.kC, b.vC, + H, KVH, dk, d.half, pos, pos*kvDim, d.posDiv, + ) + } else if e == nil { + e = firstErr( + r.RoPE(d.q.b, d.dinv.b, d.q.b, 1, D, H, dk, d.half, pos, d.posDiv), + r.RoPE(d.k.b, d.dinv.b, d.k.b, 1, kvDim, KVH, dk, d.half, pos, d.posDiv), + d.recordKVPairBlit(r, d.k.b, 0, d.v_.b, 0, b.kC, b.vC, pos*kvDim, kvDim), + ) + } + } + e = firstErr(e, recordBarrier(r)) + e = firstErr(e, d.recordMHA(r, qBuf, b.kC, b.vC, d.attn.b, 1, pos+1)) + e = firstErr(e, recordBarrier(r)) + e = firstErr(e, d.recordOProj(r, b, 1)) // dx += attn·Wo (+ optional o-bias) + e = firstErr(e, recordBarrier(r)) + e = firstErr(e, d.recordFFNSublayer(r, b, 1)) + e = firstErr(e, recordBarrier(r)) + if e != nil { + r.Free() + return nil, e + } + } + e := d.norm(r, d.dx.b, d.gFinal.b, d.bFinalBeta(), d.xn.b, 1) + e = firstErr(e, recordBarrier(r)) + e = firstErr(e, d.recordLogits(r, 1)) + if e != nil { + r.Free() + return nil, e + } + return r, nil +} + +// dropPending frees a pre-encoded next-step recorder that can no longer be used (the +// position moved unexpectedly, a prefill rewrote the cache, or the decoder is released). +func (d *Decoder) dropPending() { + if d.pending != nil { + d.pending.Free() + d.pending = nil + } + d.pendingPos = -1 +} + +// Step advances the decoder by one token at absolute position pos (== cache length before the call), +// recording the whole layer stack + vocab head into one command buffer, and returns the [vocab] +// logits. pos must be < the model's Ctx. Sequential calls run the §T614 encode-overlap: while the +// GPU executes this step, the next position's command buffer is pre-encoded on the host. +func (d *Decoder) stepInto(token, pos int) error { + if pos < 0 || pos >= d.maxLen { + return fmt.Errorf("llamagpu(%s): pos %d out of [0,%d)", d.ops.name, pos, d.maxLen) + } + if token < 0 || token >= d.v { + return fmt.Errorf("llamagpu(%s): token %d out of vocab %d", d.ops.name, token, d.v) + } + // Mamba/Mamba-2/RWKV carry per-block conv/SSM state across Step calls with no KV cache; a fresh + // sequence (pos==0) must start from zeroed state. + if pos == 0 { + if d.mamba { + if err := d.resetMambaState(); err != nil { + return err + } + } + if d.mamba2 { + if err := d.resetMamba2State(); err != nil { + return err + } + } + if d.rwkv { + if err := d.resetRWKVState(); err != nil { + return err + } + } + if d.jamba { // reset the Mamba-layer conv/SSM state (attention layers' KV cache overwrites from row 0) + if err := d.resetMambaState(); err != nil { + return err + } + } + } + // the token's embedding must be resident BEFORE commit — the recorded chain reads d.dx. + if err := d.dx.b.UploadF32(d.gatherEmbed(token)); err != nil { + return err + } + var r recorder + if d.pending != nil && d.pendingPos == pos { + r, d.pending, d.pendingPos = d.pending, nil, -1 + } else { + d.dropPending() + var err error + if r, err = d.encodeStep(pos); err != nil { + return err + } + } + if err := r.Commit(); err != nil { + r.Free() + return err + } + // encode-overlap (§T614): pre-encode pos+1 while the GPU executes pos. Failure is not fatal — the + // next Step simply encodes fresh. Disabled for Mamba: its blocks mutate resident conv/SSM state, so + // a pre-encoded pos+1 must not be recorded until pos's state writes are committed. + if d.ops.asyncEncode && !d.mamba && !d.mamba2 && !d.rwkv && pos+1 < d.maxLen { + if nr, err := d.encodeStep(pos + 1); err == nil { + d.pending, d.pendingPos = nr, pos+1 + } + } + if err := r.Wait(); err != nil { + r.Free() + return err + } + r.Free() + return nil +} + +// Step records and executes one decode step for token at pos, returning the host logits [d.v]. It wraps +// stepInto (which leaves the result in the device logits buffer) with the whole-vocab host download; the +// on-device sampling fast-path (Generate) instead reads the device logits directly via a top-k. +func (d *Decoder) Step(token, pos int) ([]float32, error) { + if err := d.stepInto(token, pos); err != nil { + return nil, err + } + out := make([]float32, d.v) + if err := d.logits.b.DownloadF32(out); err != nil { + return nil, err + } + return out, nil +} + +// StepN advances the decoder by k tokens at absolute positions pos..pos+k-1 in ONE recorded step +// (§T418): the whole layer stack runs over [k,·] rows with causal attention against the growing +// cache (row i attends up to pos+i), and the new k KV rows are appended. Returns the [k,vocab] +// logits (row i = logits after tokens[..i]). This is the prompt-PREFILL fast path — one dispatch +// round-trip for the whole prompt instead of one per token — and the target-verification step +// speculative decoding needs. pos+k must be ≤ the model's Ctx. +func (d *Decoder) StepN(tokens []int, pos int) ([]float32, error) { return d.stepN(tokens, pos, false) } + +// StepNLast is StepN but projects and downloads ONLY the final prompt row's logits (the LM head runs +// over 1 row instead of k, and the host download shrinks from k·vocab to vocab). The transformer body +// still runs all k rows to populate the KV cache, so the returned [vocab] row is bit-identical to the +// last row of StepN. Use it for prefill whose only consumer is the post-prompt token (generation) or +// which discards logits entirely (speculative/prompt-lookup KV-seeding) — the [k-1] earlier rows' LM +// head + PCIe download are pure waste there (a 512-token prompt over a 128k vocab downloads 256 MB just +// to discard all but the last row). StepN (full-k) is kept for verify callers that need every row. +func (d *Decoder) StepNLast(tokens []int, pos int) ([]float32, error) { + return d.stepN(tokens, pos, true) +} + +// stepN is StepN with an optional last-row-only mode. When lastOnly, only the FINAL prompt row's +// logits are projected by the LM head and downloaded — the transformer body still runs all k rows +// (the KV cache needs every position), but the [k-1] earlier rows' logits are never read by a caller +// that only samples the post-prompt token (Generate), so their lm_head GEMM and vocab-sized host +// download are pure waste — large for high-vocab models (a 512-token prompt over a 128k vocab downloads +// 256 MB just to discard all but the last row). It then returns a single [vocab] row. Full-k mode is +// unchanged and keeps every row for speculative/verify callers (Medusa, prompt-lookup). +func (d *Decoder) stepN(tokens []int, pos int, lastOnly bool) ([]float32, error) { + k := len(tokens) + if k == 0 { + return nil, fmt.Errorf("llamagpu(%s): StepN needs ≥1 token", d.ops.name) + } + if pos < 0 || pos+k > d.maxLen { + return nil, fmt.Errorf("llamagpu(%s): StepN [%d,%d) out of [0,%d)", d.ops.name, pos, pos+k, d.maxLen) + } + if d.mamba || d.mamba2 || d.rwkv || d.jamba { + // Mamba/Mamba-2/RWKV/Jamba's recurrence is inherently sequential (each token's state update reads the + // previous token's), so there is no batched prefill — process the prompt token by token through + // Step, which carries the recurrent state and resets it at pos==0. Returns [k, vocab]. + if lastOnly { // only the final token's logits are wanted; still step all tokens for the state + var last []float32 + for i, tok := range tokens { + row, err := d.Step(tok, pos+i) + if err != nil { + return nil, err + } + last = row + } + return append([]float32(nil), last...), nil + } + out := make([]float32, k*d.v) + for i, tok := range tokens { + row, err := d.Step(tok, pos+i) + if err != nil { + return nil, err + } + copy(out[i*d.v:(i+1)*d.v], row) + } + return out, nil + } + // a batched step rewrites the cache and scratch — any single-step encode pre-staged for + // the old position is now invalid (§T614). + d.dropPending() + host := make([]float32, k*d.d) + for i, tok := range tokens { + if tok < 0 || tok >= d.v { + return nil, fmt.Errorf("llamagpu(%s): token %d out of vocab %d", d.ops.name, tok, d.v) + } + copy(host[i*d.d:(i+1)*d.d], d.gatherEmbed(tok)) + } + if err := d.dx.b.UploadF32(host); err != nil { + return nil, err + } + r, err := d.ops.newRecorder() + if err != nil { + return nil, err + } + D, H, KVH, dk, kvDim := d.qDim, d.h, d.kvH, d.dk, d.kvDim // D = attention Q/O width (h·dk); = d.d unless head_dim decoupled + stride := D + 2*kvDim + for _, b := range d.blocks { + if d.mla { // DeepSeek-V2 MLA prefill: k new tokens at position pos + if e := firstErr(d.recordMLAAttention(r, b, pos, k), d.recordFFNSublayer(r, b, k)); e != nil { + r.Free() + return nil, e + } + continue + } + // fused QKV (§T613): one [k,stride] matmul; q/k bands rotate in place (RoPEAt's + // width parameter acts as the row stride), then Copy2D extracts q contiguously + // and deposits the k/v bands directly as cache rows. Recording order = execution + // order: norm first, then the projection branch. + e := d.recordAttnNorm(r, b, k) + attnK, attnV, attnF16 := b.kC, b.vC, d.f16KV + prefillQKV := b.wqkv + mixedQKV := prefillQKV == nil && b.wqkvPrefill != nil && k >= 24 + if mixedQKV { + prefillQKV = b.wqkvPrefill + } + if e == nil && prefillQKV != nil { + e = firstErr( + d.recordQKVProjWith(r, prefillQKV, b.qkvBias, k), + d.recordQKNorm(r, b, k, stride), // per-head QK-norm (Qwen3); no-op otherwise + ) + split := false + if sr, ok := r.(ropePairSplitRecorder); ok && mixedQKV && !d.noRope && d.rotaryDim == d.dk { + // The mixed group wins at the projection but lost its full-model gate when the + // canonical fused layout was unpacked by three standalone Copy2D dispatches. Rotate + // q/k and scatter q/k/v together; cache conversion then reads contiguous k/v. + e = firstErr(e, sr.RoPEPairSplit( + d.qkv.b, d.dinv.b, d.q.b, d.k.b, d.v_.b, + k, stride, H, 0, KVH, D, dk, d.half, D+kvDim, kvDim, pos, d.posDiv, + )) + split = true + if d.f16KV && pos == 0 { + attnK, attnV, attnF16 = d.k.b, d.v_.b, false + } + e = firstErr(e, d.recordKVPairBlit( + r, d.k.b, 0, d.v_.b, 0, b.kC, b.vC, pos*kvDim, k*kvDim, + )) + } + if !split { + e = firstErr( + e, + d.ropeQK(r, d.qkv.b, d.dinv.b, k, stride, pos), // q+k bands, ONE dispatch (§T613); full ∨ partial rotary + r.Copy2D(d.qkv.b, 0, stride, d.q.b, 0, D, k, D), + ) + // At initial prefill, retain contiguous f32 K/V scratch for the established FlashMM + // attention path while independently converting the same rows into the f16 decode cache. + if d.f16KV && pos == 0 { + e = firstErr(e, + r.Copy2D(d.qkv.b, D, stride, d.k.b, 0, kvDim, k, kvDim), + r.Copy2D(d.qkv.b, D+kvDim, stride, d.v_.b, 0, kvDim, k, kvDim), + ) + attnK, attnV, attnF16 = d.k.b, d.v_.b, false + } + e = firstErr(e, d.recordKVPairCopy2D( + r, d.qkv.b, D, stride, d.qkv.b, D+kvDim, stride, + b.kC, b.vC, pos*kvDim, kvDim, k, kvDim, + )) + } + } else if e == nil { + e = firstErr( + b.wq.record(r, d.xn.b, d.q.b, k), + b.wk.record(r, d.xn.b, d.k.b, k), + b.wv.record(r, d.xn.b, d.v_.b, k), + r.RoPE(d.q.b, d.dinv.b, d.q.b, k, D, H, dk, d.half, pos, d.posDiv), + r.RoPE(d.k.b, d.dinv.b, d.k.b, k, kvDim, KVH, dk, d.half, pos, d.posDiv), + d.recordKVPairBlit(r, d.k.b, 0, d.v_.b, 0, b.kC, b.vC, pos*kvDim, k*kvDim), + ) + if d.f16KV && pos == 0 { + attnK, attnV, attnF16 = d.k.b, d.v_.b, false + } + } + e = firstErr( + e, + // sq=k vs sk=pos+k: the kernel's causal offset (sk-sq = pos) makes row i attend + // through absolute position pos+i — exactly the prefill/verify semantics. + d.recordMHAStorage(r, d.q.b, attnK, attnV, d.attn.b, k, pos+k, attnF16), + d.recordOProj(r, b, k), // dx += attn·Wo (+ optional o-bias) + d.recordFFNSublayer(r, b, k), + ) + if e != nil { + r.Free() + return nil, e + } + } + rows := k + if lastOnly { + // Only the final row's logits are wanted: move its residual to row 0 (r.Blit copies d.d floats + // from d.dx[(k-1)·d] to d.dx[0]; a no-op when k==1) so the final norm + LM-head project ONE row + // instead of k, and the host download shrinks from k·vocab to vocab. Bit-identical: the discarded + // rows' logits were never read, and row 0's value is exactly what row k-1 would have produced. + if e := r.Blit(d.dx.b, (k-1)*d.d, d.dx.b, 0, d.d); e != nil { + r.Free() + return nil, e + } + rows = 1 + } + if e := firstErr( + d.norm(r, d.dx.b, d.gFinal.b, d.bFinalBeta(), d.xn.b, rows), + d.recordLogits(r, rows), + r.Finish(), + ); e != nil { + r.Free() + return nil, e + } + r.Free() + out := make([]float32, rows*d.v) + if err := d.logits.b.DownloadF32(out); err != nil { + return nil, err + } + return out, nil +} + +// Generate runs the batched decode as a text-generation loop: it prefills the prompt, then samples +// up to maxNew tokens (bounded by the model's Ctx), feeding each back. Returns prompt+generated +// token ids. With a greedy sampler it produces the same ids as nlp.Llama.Generate. +func (d *Decoder) Generate(prompt []int, maxNew int, s nlp.TokenSampler) ([]int, error) { + if len(prompt) == 0 { + return nil, fmt.Errorf("llamagpu(%s): Generate needs a non-empty prompt", d.ops.name) + } + out := append([]int(nil), prompt...) + // prefill the whole prompt in ONE recorded step (§T418) — one dispatch round-trip instead of one + // per prompt token. lastOnly: generation samples only the post-prompt row, so the LM head projects + // and downloads that single row (not all len(prompt) rows) — a large saving for high-vocab models. + all, err := d.stepN(prompt, 0, true) + if err != nil { + return nil, err + } + pos := len(prompt) + logits := all // last-row-only prefill returns exactly the post-prompt row's [vocab] logits + buf := make([]float64, d.v) + // Device-resident top-k sampling fast-path: for a penalty-free TopK sampler AND a logits buffer that + // can TopK (CUDA GPU reduction or Metal coherent-UMA selection), skip the whole-vocab host download. + // Read the K highest logits at the resident boundary and draw over them, bit-identical to the + // full-vocab sampler (see fastTopKSampler/sampleTopKCandidates). Any other sampler/backend takes the + // flexible host path. + sp, fastK := fastTopKSampler(s, d.v) + spP, fastC := fastTopPSampler(s, d.v) // pure-top-p device path (mutually exclusive with sp: needs top-k off) + dt, hasDev := d.logits.b.(deviceTopKer) + dtP, hasDevP := d.logits.b.(deviceTopPer) + fast := (sp != nil && hasDev) || (spP != nil && hasDevP) + for range maxNew { + if pos >= d.maxLen { + break + } + var next int + if fast && logits == nil { // decode step, logits resident on device (first d.v of the buffer) + if sp != nil { // top-k (± post-shrink top-p/min-p): draw over the K-candidate superset + gi, gv, e := dt.TopKN(d.v, fastK) + if e != nil { + return nil, e + } + next = sampleTopKCandidates(sp, gi, gv) + } else { // pure top-p: resolve the nucleus over the top-C candidates using full-vocab stats + maxL, zexp, e2 := dtP.SoftmaxStatsN(d.v, spP.Temperature) + if e2 != nil { + return nil, e2 + } + resolved := false + // Cheap overflow guard (stats only): top-C mass ≤ C/Zexp; below TopP the nucleus can't fit, + // so skip the O(n·C) device TopK and fall straight to the host path (see the graph decoder). + if float64(fastC)/zexp >= spP.TopP { + gi, gv, e := dtP.TopKN(d.v, fastC) + if e != nil { + return nil, e + } + cl := make([]float64, len(gv)) + for i, v := range gv { + cl[i] = float64(v) + } + if tok, ok := spP.SampleTopPFromCandidates(cl, gi, maxL, zexp); ok { + next = tok + resolved = true + } + } + if !resolved { // guard predicted overflow, or TopK confirmed it → full-vocab host sample + l, e3 := dtP.ToHost() + if e3 != nil { + return nil, e3 + } + lf := l.Storage().F32() + for i := 0; i < d.v; i++ { + buf[i] = float64(lf[i]) + } + next = s.SampleWithHistory(buf, out) + } + } + } else { // first token (prefill host logits) or the full-vocab fallback + for i, x := range logits { + buf[i] = float64(x) + } + next = s.SampleWithHistory(buf, out) + } + out = append(out, next) + if fast { + if e := d.stepInto(next, pos); e != nil { // leaves logits on device — no host download + return nil, e + } + logits = nil + } else { + l, e := d.Step(next, pos) + if e != nil { + return nil, e + } + logits = l + } + pos++ + } + return out, nil +} + +// Vocab returns the model's vocabulary size. +func (d *Decoder) Vocab() int { return d.v } + +// Ctx returns the model's maximum context length (the KV-cache capacity). +func (d *Decoder) Ctx() int { return d.maxLen } + +// Release frees all device buffers and resident quantized weights. +func (d *Decoder) Release() { + d.dropPending() + for _, b := range d.all { + if b != nil { + b.Release() + } + } + d.all = nil + for _, w := range d.qweights { + if w != nil { + _ = w.Close() + } + } + d.qweights = nil +} + +func boolToInt(b bool) int { + if b { + return 1 + } + return 0 +} + +func firstErr(errs ...error) error { + for _, e := range errs { + if e != nil { + return e + } + } + return nil +} + +// rowBlock is the number of source rows converted at a time by forEachRowBlockF32; colTile is the +// column tile flat2DT transposes within such a block. Together they make the transpose a classic +// cache-blocked 2D tile walk: both the contiguous reads and the strided writes stay inside a +// cache-resident working set instead of touching a fresh line (and TLB entry) per element. The +// pair was picked by measurement on a [32000, 2048] lm_head — see the §V22 benchmarks in +// flatten_test.go, where tiling is worth ~1.8× on top of the bulk-slice conversion alone. +const ( + rowBlock = 256 + colTile = 64 +) + +// forEachRowBlockF32 walks t (rank-2, any dtype, any strides) in blocks of at most rowBlock rows, +// calling fn with the block's first row index and its values converted to float32 in contiguous +// row-major order (len(blk) == (j1-j0)·cols). +// +// This is the bulk-slice idiom the base-perf sweep established: one dtype decision per BLOCK via +// tensor.Cast's typed fast paths, instead of one AtF64 interface dispatch per ELEMENT. It is +// numerically identical to float32(t.AtF64(j, i)) for every dtype — Cast's F64→F32, F16→F32 and +// BF16→F32 paths reproduce the widen-then-narrow composition exactly — so callers stay +// bit-identical to the per-element loops they replace. +func forEachRowBlockF32(t *tensor.Tensor, fn func(j0 int, blk []float32)) { + rows, cols := t.Shape()[0], t.Shape()[1] + //perfscan:ignore PS2001 Cast alloc per row-BLOCK (256), model-build bulk-slice upload, amortized + for j0 := 0; j0 < rows; j0 += rowBlock { + j1 := min(j0+rowBlock, rows) + blk, err := t.Slice(0, j0, j1) + if err != nil { // unreachable for in-range slices; degrade to the per-element read + scratch := make([]float32, (j1-j0)*cols) + for j := j0; j < j1; j++ { + for i := range cols { + scratch[(j-j0)*cols+i] = float32(t.AtF64(j, i)) + } + } + fn(j0, scratch) + continue + } + if blk.Dtype() == tensor.F32 && blk.IsContiguous() { + // Already the target layout and dtype: hand out the storage window itself, no + // conversion and no allocation. This is the f32-weight upload path. + off := blk.Offset() + fn(j0, blk.Storage().F32()[off:off+(j1-j0)*cols]) + continue + } + fn(j0, blk.Cast(tensor.F32).Storage().F32()[:(j1-j0)*cols]) + } +} + +// flat2DT flattens the TRANSPOSE of a rank-2 tensor into a row-major float32 slice: +// out[i·rows+j] == float32(t.AtF64(j, i)) for t of shape [rows, cols], i.e. the [cols, rows] +// layout the device matmuls expect from a torch-style [out, in] weight or a tied lm_head. +// +// Used for weight upload at model-build time, where the tensor can be huge — Gemma's tied lm_head +// is [256000, 2048] = 524M elements — so this runs as a blocked bulk-slice transpose rather than +// per-element reads. +func flat2DT(t *tensor.Tensor) []float32 { + o := make([]float32, t.Shape()[0]*t.Shape()[1]) + flat2DTInto(o, t) + return o +} + +// flat2DTInto is flat2DT writing into a caller-provided destination of exactly rows·cols floats. +// It lets a caller that already owns the target buffer — e.g. the tied-lm_head constructors, which +// need the result inside a [dim, vocab] tensor — skip an intermediate slice, which at Gemma's +// 256k vocab is a 2 GB allocation and a 2 GB copy saved. +func flat2DTInto(o []float32, t *tensor.Tensor) { + rows, cols := t.Shape()[0], t.Shape()[1] + forEachRowBlockF32(t, func(j0 int, blk []float32) { + j1 := j0 + len(blk)/cols + for i0 := 0; i0 < cols; i0 += colTile { + i1 := min(i0+colTile, cols) + for j := j0; j < j1; j++ { + row := blk[(j-j0)*cols+i0 : (j-j0)*cols+i1] + for i, v := range row { + o[(i0+i)*rows+j] = v + } + } + } + }) +} + +// flat2D flattens a rank-2 tensor into a row-major float32 slice (no transpose): +// out[i·cols+j] == float32(t.AtF64(i, j)). +func flat2D(t *tensor.Tensor) []float32 { + r, c := t.Shape()[0], t.Shape()[1] + o := make([]float32, r*c) + forEachRowBlockF32(t, func(j0 int, blk []float32) { copy(o[j0*c:], blk) }) + return o +} + +// flat1D flattens a rank-1 tensor (a norm γ, a bias) into a float32 slice. +func flat1D(t *tensor.Tensor) []float32 { + n := t.Shape()[0] + o := make([]float32, n) + copy(o, t.Cast(tensor.F32).Storage().F32()[:n]) + return o +} + +// embedRow reads one row of an embedding table as float32. It is on the per-token decode path, so +// it takes the bulk-slice route too: one row view, one typed conversion, no per-element dispatch. +// +//perfscan:ignore PS3033 one small row alloc/token, resource-only, negligible vs forward +func embedRow(table *tensor.Tensor, row, cols int) []float32 { + o := make([]float32, cols) + r, err := table.Slice(0, row, row+1) + if err != nil { + for j := range cols { + o[j] = float32(table.AtF64(row, j)) + } + return o + } + copy(o, r.Cast(tensor.F32).Storage().F32()[:cols]) + return o +} + +// gatherEmbed reads a token's embedding row and applies the Granite EmbeddingMult (embMult == 1 +// for every non-Granite model, so this is embedRow verbatim then). +func (d *Decoder) gatherEmbed(token int) []float32 { + e := embedRow(d.table, token, d.d) + if d.embMult != 1 { + for i := range e { + e[i] *= d.embMult + } + } + return e +} diff --git a/checks/testdata/ps6136-owner/before/gpt.go.txt b/checks/testdata/ps6136-owner/before/gpt.go.txt new file mode 100644 index 00000000..316d4b17 --- /dev/null +++ b/checks/testdata/ps6136-owner/before/gpt.go.txt @@ -0,0 +1,389 @@ +package llamagpu + +import ( + "fmt" + "math" + + "github.com/jxsl13/goai/nlp" + "github.com/jxsl13/goai/tensor" +) + +// gptBlock is one pre-LN GPT transformer block on the device: LayerNorm gains/biases, the four +// attention projections, the biased GELU MLP, and the block's KV cache. +type gptBlock struct { + g1, b1, g2, b2 buffer // LN1/LN2 gamma+beta [d] + wq, wk, wv, wo linear // attention projections [d,d], no bias + wqkv f32Linear // fused attention projection [d,3d] + w1, w2 linear // FFN up [d,4d] / down [4d,d] + fb1, fb2 buffer // FFN biases [4d] / [d] + kC, vC buffer // KV caches [maxLen*d] + ffn int // FFN width, fixed when the block is uploaded +} + +// GPTDecoder is the batched-decode engine for nlp.GPT models (§T422) — the GPT-2-style sibling of +// [Decoder]: LayerNorm instead of RMSNorm, learned positional embeddings instead of RoPE, a biased +// GELU MLP instead of SwiGLU. Each Step records the whole layer stack into one command buffer. +// Not safe for concurrent use; Release when done. +type GPTDecoder struct { + ops backendOps + d, h, dk, v, maxLen int + eps, scale float32 + + blocks []gptBlock + lnfG, lnfB *bufSlot + head linear + dx, xn, xn2 *bufSlot + q, k, v_, qkv *bufSlot + attn, hid *bufSlot + logits *bufSlot + tokEmb, posEmb *tensor.Tensor // host gather sources + all []buffer +} + +type f32QKVBandsRecorder interface { + F32QKVBands(x, w, q, k, v buffer, rows, dim int) error +} + +type biasGELURecorder interface { + BiasGELU(x, b, o buffer, rows, n int) error +} + +func (d *GPTDecoder) recordAttentionResidual(r recorder, b gptBlock, rows int) error { + return b.wo.recordAdd(r, d.attn.b, nil, d.dx.b, rows, d.d) +} + +func (d *GPTDecoder) recordFFNResidual(r recorder, b gptBlock, rows int) error { + return firstErr( + b.w2.recordAdd(r, d.hid.b, nil, d.dx.b, rows, d.d), + r.AddBias(d.dx.b, b.fb2, d.dx.b, rows, d.d), + ) +} + +func (d *GPTDecoder) recordBiasGELU(r recorder, b gptBlock, rows int) error { + if d.ops.fusedBiasGELU { + if fused, ok := r.(biasGELURecorder); ok { + return fused.BiasGELU(d.hid.b, b.fb1, d.hid.b, rows, b.ffn) + } + } + return firstErr( + r.AddBias(d.hid.b, b.fb1, d.hid.b, rows, b.ffn), + r.Unary(d.hid.b, d.hid.b, unaryGELU), + ) +} + +// newGPTDecoder uploads m's weights via ops and prepares per-block KV caches up to Ctx. +func newGPTDecoder(m *nlp.GPT, ops backendOps) (*GPTDecoder, error) { + cfg := m.Config + d := &GPTDecoder{ + ops: ops, + d: cfg.Dim, h: cfg.Heads, v: cfg.Vocab, maxLen: cfg.Ctx, + eps: float32(cfg.Eps), tokEmb: m.TokEmb, posEmb: m.PosEmb, + } + if d.h <= 0 || d.d%d.h != 0 { + return nil, fmt.Errorf("llamagpu(%s): gpt dim %d not divisible by heads %d", ops.name, d.d, d.h) + } + d.dk = d.d / d.h + d.scale = float32(1.0 / math.Sqrt(float64(d.dk))) + + var err error + mk := func(data []float32) *bufSlot { + if err != nil { + return &bufSlot{} + } + b, e := ops.newBuffer(data) + if e != nil { + err = e + return &bufSlot{} + } + d.all = append(d.all, b) + return &bufSlot{b: b} + } + lin := func(w *tensor.Tensor) linear { + in, out := w.Shape()[0], w.Shape()[1] + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } + fusedQKV := func(wq, wk, wv *tensor.Tensor) f32Linear { + in := wq.Shape()[0] + nq, nk, nv := wq.Shape()[1], wk.Shape()[1], wv.Shape()[1] + fq, fk, fv := flat2D(wq), flat2D(wk), flat2D(wv) + stride := nq + nk + nv + weight := make([]float32, in*stride) + for row := range in { + dst := weight[row*stride : (row+1)*stride] + copy(dst[:nq], fq[row*nq:(row+1)*nq]) + copy(dst[nq:nq+nk], fk[row*nk:(row+1)*nk]) + copy(dst[nq+nk:], fv[row*nv:(row+1)*nv]) + } + return f32Linear{w: mk(weight).b, k: in, n: stride} + } + for _, b := range m.Blocks { + gb := gptBlock{ + g1: mk(flat1D(b.LN1.Gamma)).b, b1: mk(flat1D(b.LN1.Beta)).b, + g2: mk(flat1D(b.LN2.Gamma)).b, b2: mk(flat1D(b.LN2.Beta)).b, + wo: lin(b.Attn.Wo), w1: lin(b.W1), w2: lin(b.W2), + fb1: mk(flat1D(b.B1)).b, fb2: mk(flat1D(b.B2)).b, + kC: mk(make([]float32, d.maxLen*d.d)).b, vC: mk(make([]float32, d.maxLen*d.d)).b, + ffn: b.W1.Shape()[1], + } + if ops.fusedF32QKV { + gb.wqkv = fusedQKV(b.Attn.Wq, b.Attn.Wk, b.Attn.Wv) + } else { + gb.wq, gb.wk, gb.wv = lin(b.Attn.Wq), lin(b.Attn.Wk), lin(b.Attn.Wv) + } + d.blocks = append(d.blocks, gb) + } + d.lnfG = mk(flat1D(m.LNf.Gamma)) + d.lnfB = mk(flat1D(m.LNf.Beta)) + d.head = lin(m.Head) + ffn := 4 * d.d + if len(m.Blocks) > 0 { + ffn = m.Blocks[0].W1.Shape()[1] + } + c := d.maxLen // scratch sized for multi-token StepN (§T423) + d.dx = mk(make([]float32, c*d.d)) + d.xn = mk(make([]float32, c*d.d)) + d.xn2 = mk(make([]float32, c*d.d)) + d.q = mk(make([]float32, c*d.d)) + d.k = mk(make([]float32, c*d.d)) + d.v_ = mk(make([]float32, c*d.d)) + if ops.fusedF32QKV { + // Large prefills address bands of the single fused weight directly. Only rows below the + // strided-band crossover need a grouped-output scratch, bounding it to 63 rows. + d.qkv = mk(make([]float32, min(c, 63)*3*d.d)) + } + d.attn = mk(make([]float32, c*d.d)) + d.hid = mk(make([]float32, c*ffn)) + d.logits = mk(make([]float32, c*d.v)) + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// Step advances the decoder by one token at absolute position pos (== cache length before the +// call), recording the whole layer stack + head into one command buffer, and returns the [vocab] +// logits. The token embedding and learned positional embedding are summed host-side. +func (d *GPTDecoder) Step(token, pos int) ([]float32, error) { + if pos < 0 || pos >= d.maxLen { + return nil, fmt.Errorf("llamagpu(%s): gpt pos %d out of [0,%d)", d.ops.name, pos, d.maxLen) + } + if token < 0 || token >= d.v { + return nil, fmt.Errorf("llamagpu(%s): gpt token %d out of vocab %d", d.ops.name, token, d.v) + } + x := embedRow(d.tokEmb, token, d.d) + pe := embedRow(d.posEmb, pos, d.d) + for i := range x { + x[i] += pe[i] + } + if err := d.dx.b.UploadF32(x); err != nil { + return nil, err + } + r, err := d.ops.newRecorder() + if err != nil { + return nil, err + } + D, H, dk := d.d, d.h, d.dk + for _, b := range d.blocks { + e := r.LayerNorm(d.dx.b, b.g1, b.b1, d.xn.b, 1, D, d.eps) + qBuf := d.q.b + if d.ops.fusedF32QKV { + qBuf = d.qkv.b + e = firstErr(e, + b.wqkv.record(r, d.xn.b, d.qkv.b, 1), + r.Blit(d.qkv.b, D, b.kC, pos*D, D), + r.Blit(d.qkv.b, 2*D, b.vC, pos*D, D), + ) + } else { + e = firstErr(e, + b.wq.record(r, d.xn.b, d.q.b, 1), + b.wk.record(r, d.xn.b, d.k.b, 1), + b.wv.record(r, d.xn.b, d.v_.b, 1), + r.Blit(d.k.b, 0, b.kC, pos*D, D), + r.Blit(d.v_.b, 0, b.vC, pos*D, D), + ) + } + e = firstErr(e, + r.MHA(qBuf, b.kC, b.vC, d.attn.b, 1, pos+1, D, H, H, dk, 1, 0, d.scale), + d.recordAttentionResidual(r, b, 1), + r.LayerNorm(d.dx.b, b.g2, b.b2, d.xn2.b, 1, D, d.eps), + b.w1.record(r, d.xn2.b, d.hid.b, 1), + d.recordBiasGELU(r, b, 1), + d.recordFFNResidual(r, b, 1), + ) + if e != nil { + r.Free() + return nil, e + } + } + if e := firstErr( + r.LayerNorm(d.dx.b, d.lnfG.b, d.lnfB.b, d.xn.b, 1, D, d.eps), + d.head.record(r, d.xn.b, d.logits.b, 1), + r.Finish(), + ); e != nil { + r.Free() + return nil, e + } + r.Free() + out := make([]float32, d.v) + if err := d.logits.b.DownloadF32(out); err != nil { + return nil, err + } + return out, nil +} + +// StepN advances the decoder by k tokens at absolute positions pos..pos+k-1 in ONE recorded step +// (§T423, the GPT sibling of Decoder.StepN): the layer stack runs over [k,·] rows with causal +// attention against the growing cache, biases broadcast via the AddBias record op, and the k KV +// rows are appended. Returns the [k,vocab] logits. +func (d *GPTDecoder) StepN(tokens []int, pos int) ([]float32, error) { + return d.gptStepN(tokens, pos, false) +} + +// StepNLast is StepN projecting only the FINAL prompt row's logits (LM head over 1 row, download of +// vocab not k·vocab); the transformer body still runs all k rows so the KV cache and the returned +// [vocab] row are bit-identical to StepN's last row. For generation/KV-seeding prefills that read at +// most the last row — mirrors Decoder.StepNLast. +func (d *GPTDecoder) StepNLast(tokens []int, pos int) ([]float32, error) { + return d.gptStepN(tokens, pos, true) +} + +func (d *GPTDecoder) gptStepN(tokens []int, pos int, lastOnly bool) ([]float32, error) { + k := len(tokens) + if k == 0 { + return nil, fmt.Errorf("llamagpu(%s): gpt StepN needs ≥1 token", d.ops.name) + } + if pos < 0 || pos+k > d.maxLen { + return nil, fmt.Errorf("llamagpu(%s): gpt StepN [%d,%d) out of [0,%d)", d.ops.name, pos, pos+k, d.maxLen) + } + host := make([]float32, k*d.d) + for i, tok := range tokens { + if tok < 0 || tok >= d.v { + return nil, fmt.Errorf("llamagpu(%s): gpt token %d out of vocab %d", d.ops.name, tok, d.v) + } + row := host[i*d.d : (i+1)*d.d] + copy(row, embedRow(d.tokEmb, tok, d.d)) + pe := embedRow(d.posEmb, pos+i, d.d) + for j := range row { + row[j] += pe[j] + } + } + if err := d.dx.b.UploadF32(host); err != nil { + return nil, err + } + r, err := d.ops.newRecorder() + if err != nil { + return nil, err + } + D, H, dk := d.d, d.h, d.dk + for _, b := range d.blocks { + e := r.LayerNorm(d.dx.b, b.g1, b.b1, d.xn.b, k, D, d.eps) + if d.ops.fusedF32QKV { + if bands, ok := r.(f32QKVBandsRecorder); ok && k >= 64 { + e = firstErr(e, bands.F32QKVBands(d.xn.b, b.wqkv.w, d.q.b, d.k.b, d.v_.b, k, D)) + } else { + e = firstErr(e, + b.wqkv.record(r, d.xn.b, d.qkv.b, k), + r.Copy2D(d.qkv.b, 0, 3*D, d.q.b, 0, D, k, D), + r.Copy2D(d.qkv.b, D, 3*D, d.k.b, 0, D, k, D), + r.Copy2D(d.qkv.b, 2*D, 3*D, d.v_.b, 0, D, k, D), + ) + } + } else { + e = firstErr(e, + b.wq.record(r, d.xn.b, d.q.b, k), + b.wk.record(r, d.xn.b, d.k.b, k), + b.wv.record(r, d.xn.b, d.v_.b, k), + ) + } + e = firstErr(e, + r.Blit(d.k.b, 0, b.kC, pos*D, k*D), + r.Blit(d.v_.b, 0, b.vC, pos*D, k*D), + r.MHA(d.q.b, b.kC, b.vC, d.attn.b, k, pos+k, D, H, H, dk, 1, 0, d.scale), + d.recordAttentionResidual(r, b, k), + r.LayerNorm(d.dx.b, b.g2, b.b2, d.xn2.b, k, D, d.eps), + b.w1.record(r, d.xn2.b, d.hid.b, k), + d.recordBiasGELU(r, b, k), + d.recordFFNResidual(r, b, k), + ) + if e != nil { + r.Free() + return nil, e + } + } + rows := k + if lastOnly { + // Only the final row's logits are wanted: move its residual to row 0 so the final LN + LM head + // project ONE row and the download shrinks from k·vocab to vocab (bit-identical — row 0 then + // holds exactly what row k-1 would have produced; a no-op when k==1). + if e := r.Blit(d.dx.b, (k-1)*D, d.dx.b, 0, D); e != nil { + r.Free() + return nil, e + } + rows = 1 + } + if e := firstErr( + r.LayerNorm(d.dx.b, d.lnfG.b, d.lnfB.b, d.xn.b, rows, D, d.eps), + d.head.record(r, d.xn.b, d.logits.b, rows), + r.Finish(), + ); e != nil { + r.Free() + return nil, e + } + r.Free() + out := make([]float32, rows*d.v) + if err := d.logits.b.DownloadF32(out); err != nil { + return nil, err + } + return out, nil +} + +// Generate prefills the whole prompt in ONE recorded StepN (§T423), then samples up to maxNew +// tokens (bounded by Ctx), returning prompt+generated ids. +func (d *GPTDecoder) Generate(prompt []int, maxNew int, s nlp.TokenSampler) ([]int, error) { + if len(prompt) == 0 { + return nil, fmt.Errorf("llamagpu(%s): gpt Generate needs a non-empty prompt", d.ops.name) + } + out := append([]int(nil), prompt...) + // generation samples only the post-prompt row, so project just the last prompt row's logits + // (StepNLast) instead of the full [len(prompt), vocab] LM head + download. + logits, err := d.StepNLast(prompt, 0) + if err != nil { + return nil, err + } + pos := len(prompt) + buf := make([]float64, d.v) + for range maxNew { + if pos >= d.maxLen { + break + } + for i, x := range logits { + buf[i] = float64(x) + } + next := s.SampleWithHistory(buf, out) + out = append(out, next) + l, err := d.Step(next, pos) + if err != nil { + return nil, err + } + logits = l + pos++ + } + return out, nil +} + +// Vocab returns the model's vocabulary size. +func (d *GPTDecoder) Vocab() int { return d.v } + +// Ctx returns the model's maximum context length (the KV-cache capacity). +func (d *GPTDecoder) Ctx() int { return d.maxLen } + +// Release frees all device buffers. +func (d *GPTDecoder) Release() { + for _, b := range d.all { + if b != nil { + b.Release() + } + } + d.all = nil +} diff --git a/checks/testdata/ps6136-owner/cuda.go.txt b/checks/testdata/ps6136-owner/cuda.go.txt new file mode 100644 index 00000000..771ef96e --- /dev/null +++ b/checks/testdata/ps6136-owner/cuda.go.txt @@ -0,0 +1,1504 @@ +//go:build cuda && cgo && (linux || windows) + +package llamagpu + +import ( + "fmt" + + "github.com/jxsl13/goai/backend" + "github.com/jxsl13/goai/backend/cuda" + "github.com/jxsl13/goai/format/gguf" + "github.com/jxsl13/goai/nlp" + "github.com/jxsl13/goai/tensor" +) + +// cuda adapter: thin assertions from the backend-agnostic buffer/recorder interfaces to the +// concrete cuda.DeviceF32/cuda.Recorder, the same shape as the metal (§T408) and vulkan (§T409) +// adapters. The one difference is submit semantics: the CUDA recorder enqueues eagerly on the work +// stream (Commit is a no-op, Wait syncs), so asyncEncode stays false — a second recorder must not +// pre-encode onto the shared stream while one is in flight. + +type cBuf struct{ *cuda.DeviceF32 } + +type cRec struct{ r *cuda.Recorder } + +func (c cRec) RMSNorm(x, g, o buffer, rows, dim int, eps float32) error { + return c.r.RMSNorm(cb(x), cb(g), cb(o), rows, dim, eps) +} +func (c cRec) LayerNorm(x, g, b, o buffer, rows, dim int, eps float32) error { + return c.r.LayerNorm(cb(x), cb(g), cb(b), cb(o), rows, dim, eps) +} +func (c cRec) AddBias(x, b, o buffer, rows, n int) error { + return c.r.AddBias(cb(x), cb(b), cb(o), rows, n) +} +func (c cRec) MatMul(a, b, cc buffer, m, k, n int) error { + return c.r.MatMul(cb(a), cb(b), cb(cc), m, k, n) +} +func (c cRec) MatMulAcc(a, b, cc buffer, m, k, n int) error { + return c.r.MatMulAcc(cb(a), cb(b), cb(cc), m, k, n) +} +func (c cRec) RoPE(q, inv, o buffer, seq, width, heads, hd, half, pos int, posDiv float32) error { + return c.r.RoPE(cb(q), cb(inv), cb(o), seq, width, heads, hd, half, pos, posDiv) +} +func (c cRec) RoPEAt(q, inv, o buffer, off, seq, width, heads, hd, half, pos int, posDiv float32) error { + return c.r.RoPEAt(cb(q), cb(inv), cb(o), off, seq, width, heads, hd, half, pos, posDiv) +} +func (c cRec) RoPEPair(qkv, inv buffer, seq, stride, headsQ, offQ, headsK, offK, hd, half, pos int, posDiv float32) error { + return c.r.RoPEPair(cb(qkv), cb(inv), seq, stride, headsQ, offQ, headsK, offK, hd, half, pos, posDiv) +} +func (c cRec) RoPEPartialPair(qkv, inv buffer, seq, stride, headsQ, offQ, headsK, offK, hd, rotaryDim, pos int, posDiv float32) error { + return c.r.RoPEPartialPair(cb(qkv), cb(inv), seq, stride, headsQ, offQ, headsK, offK, hd, rotaryDim, pos, posDiv) +} +func (c cRec) Blit(src buffer, srcOff int, dst buffer, dstOff, n int) error { + return c.r.Blit(cb(src), srcOff, cb(dst), dstOff, n) +} +func (c cRec) Copy2D(src buffer, srcOff, srcStride int, dst buffer, dstOff, dstStride, rows, rowFloats int) error { + return c.r.Copy2D(cb(src), srcOff, srcStride, cb(dst), dstOff, dstStride, rows, rowFloats) +} +func (c cRec) MHA(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error { + return c.r.MHA(cb(q), cb(k), cb(v), cb(o), sq, sk, dm, heads, kvHeads, dk, causal, window, scale) +} +func (c cRec) MHACap(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale, cap float32) error { + return c.r.MHACap(cb(q), cb(k), cb(v), cb(o), sq, sk, dm, heads, kvHeads, dk, causal, window, scale, cap) +} +func (c cRec) MHAALiBi(q, k, v, o, slopes buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error { + return c.r.MHAALiBi(cb(q), cb(k), cb(v), cb(o), cb(slopes), sq, sk, dm, heads, kvHeads, dk, causal, window, scale) +} +func (c cRec) MHABias(q, k, v, o, bias buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error { + return c.r.MHABias(cb(q), cb(k), cb(v), cb(o), cb(bias), sq, sk, dm, heads, kvHeads, dk, causal, window, scale) +} +func (c cRec) MoEGate(logits, weights buffer, rows, e, k, raw int, scale float32) error { + return c.r.MoEGate(cb(logits), cb(weights), rows, e, k, raw, scale) +} +func (c cRec) RowAxpy(dst, src, arow buffer, rows, cols int) error { + return c.r.RowAxpy(cb(dst), cb(src), cb(arow), rows, cols) +} +func (c cRec) MHARect(q, k, v, o buffer, sq, sk, heads, kvHeads, dqk, dv, causal, window int, scale float32) error { + return c.r.MHARect(cb(q), cb(k), cb(v), cb(o), sq, sk, heads, kvHeads, dqk, dv, causal, window, scale) +} +func (c cRec) SSMStep(u, delta, a, b, cc, dskip, h, y buffer, d, n int) error { + return c.r.SSMStep(cb(u), cb(delta), cb(a), cb(b), cb(cc), cb(dskip), cb(h), cb(y), d, n) +} +func (c cRec) Conv1DStep(x, w, b, state, out buffer, d, k int) error { + return c.r.Conv1DStep(cb(x), cb(w), cb(b), cb(state), cb(out), d, k) +} +func (c cRec) SSDStep(x, delta, a, b, cc, dskip, state, y buffer, heads, headDim, groups, n int) error { + return c.r.SSDStep(cb(x), cb(delta), cb(a), cb(b), cb(cc), cb(dskip), cb(state), cb(y), heads, headDim, groups, n) +} +func (c cRec) WKVStep(k, v, w, u, aa, bb, pp, out buffer, d int) error { + return c.r.WKVStep(cb(k), cb(v), cb(w), cb(u), cb(aa), cb(bb), cb(pp), cb(out), d) +} +func (c cRec) Unary(x, o buffer, op int) error { return c.r.Unary(cb(x), cb(o), op) } +func (c cRec) Binary(a, b, o buffer, op int) error { + return c.r.Binary(cb(a), cb(b), cb(o), op) +} + +// SwiGLUHalves fuses SwiGLU over the two halves of a column-fused gate|up buffer (§gate-up fusion): +// out[r,i] = silu(gu[r,i])·gu[r,hidden+i]. Satisfies the decoder's optional swiGLUHalvesRecorder so +// recordFFN can fuse ffn_gate|ffn_up into one GEMV + this op instead of two GEMVs + a Binary SwiGLU. +func (c cRec) SwiGLUHalves(gu, out buffer, rows, hidden int) error { + return c.r.SwiGLUHalves(cb(gu), cb(out), rows, hidden) +} + +// GeGLUHalves fuses GeGLU over the halves of a column-fused gate|up buffer (Gemma): out[r,i] = +// GELU(gu[r,i])·gu[r,hidden+i]. Satisfies the decoder's optional geGLUHalvesRecorder. +func (c cRec) GeGLUHalves(gu, out buffer, rows, hidden int) error { + return c.r.GeGLUHalves(cb(gu), cb(out), rows, hidden) +} + +// QMatMulResidentMoE routes a Q8 expert GEMV through the moeW-gated kernel (skips non-selected tokens' +// weight-stream, bit-identical to dense+RowAxpy-w0). Non-Q8 quant → dense QMatMulResident (correct). +func (c cRec) QMatMulResidentMoE(x buffer, w qweight, o buffer, m int, moeGate buffer, gateStride, ex int) error { + switch rw := w.(type) { + case *cuda.ResidentBQ8: + return rw.QMatMulMoeInto(cb(x), cb(o), cb(moeGate), m, gateStride, ex) + case *cuda.ResidentBQ4K: + return rw.QMatMulMoeInto(cb(x), cb(o), cb(moeGate), m, gateStride, ex) + } + return c.QMatMulResident(x, w, o, m) +} + +func (c cRec) QMatMulResident(x buffer, w qweight, o buffer, m int) error { + switch rw := w.(type) { + case *cuda.ResidentBQ8: + return c.r.QMatMulResident(cb(x), rw, cb(o), m) // Q8: GEMV decode + int8-tensor-core MMQ prefill + case *cuda.ResidentBQ4K: + return c.r.QMatMulResidentQ4K(cb(x), rw, cb(o), m) // Q4_K: aggressive-quant decode GEMV + tensor-core prefill + case *cuda.ResidentBQ6K: + return c.r.QMatMulResidentQ6K(cb(x), rw, cb(o), m) // Q6_K: 6-bit k-quant decode GEMV + tensor-core prefill + case *cuda.ResidentBQ5K: + return c.r.QMatMulResidentQ5K(cb(x), rw, cb(o), m) // Q5_K: 5-bit k-quant decode GEMV + tensor-core prefill + case *cuda.ResidentBQ2K: + return c.r.QMatMulResidentQ2K(cb(x), rw, cb(o), m) // Q2_K: 2-bit k-quant (0.25 B/w) decode GEMV + tensor-core prefill + case *cuda.ResidentBQ3K: + return c.r.QMatMulResidentQ3K(cb(x), rw, cb(o), m) // Q3_K: 3-bit k-quant decode GEMV + tensor-core prefill + default: + return fmt.Errorf("llamagpu: cuda QMatMulResident: unsupported resident weight %T", w) + } +} + +// QMatMulResidentAcc fuses the residual add into the resident quant GEMV (dst += x·dequant(w), beta=1) +// for the decode epilogue. All K-quants + Q8 are wired (Q8/Q4_K/Q6_K/Q5_K/Q2_K/Q3_K — the mixes +// spread projections across these); other formats return errQuantAccUnsupported → GEMV+Binary fallback. +func (c cRec) QMatMulResidentAcc(x buffer, w qweight, dst buffer, m int) error { + switch rw := w.(type) { + case *cuda.ResidentBQ8: + return c.r.QMatMulResidentAcc(cb(x), rw, cb(dst), m) + case *cuda.ResidentBQ4K: + return c.r.QMatMulResidentAccQ4K(cb(x), rw, cb(dst), m) + case *cuda.ResidentBQ6K: + return c.r.QMatMulResidentAccQ6K(cb(x), rw, cb(dst), m) + case *cuda.ResidentBQ5K: + return c.r.QMatMulResidentAccQ5K(cb(x), rw, cb(dst), m) + case *cuda.ResidentBQ2K: + return c.r.QMatMulResidentAccQ2K(cb(x), rw, cb(dst), m) + case *cuda.ResidentBQ3K: + return c.r.QMatMulResidentAccQ3K(cb(x), rw, cb(dst), m) + } + return errQuantAccUnsupported +} + +func (c cRec) Commit() error { return c.r.Commit() } +func (c cRec) Wait() error { return c.r.Wait() } +func (c cRec) Finish() error { return c.r.Finish() } +func (c cRec) Free() { c.r.Free() } + +func cb(b buffer) *cuda.DeviceF32 { return b.(cBuf).DeviceF32 } + +// NewCUDA uploads an nlp.Llama's weights into CUDA device buffers and prepares a KV cache up to +// m.Config.Ctx tokens — the NVIDIA variant of New/NewVulkan. It reuses the backend-agnostic batched +// Decoder core (§T409) via the cuda.Recorder, so Generate/Step/StepN and every sampler work +// unchanged. f32 path (quantized NewQuantCUDA is a follow-up: it needs cuda.Recorder.QMatMulResident). +func NewCUDA(m *nlp.Llama) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, // one shared work stream — pre-encoding a second recorder would interleave + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewLlamaQ8CUDA is NewCUDA with every projection (fused QKV, o_proj, gate/up/down, lm_head) quantized +// to resident Q8_0 — a direct f32→Q8 decode path that needs NO ggml QuantLlama conversion. Transformer +// decode is weight-bandwidth-bound, so streaming ~4× smaller Q8 projections speeds it up substantially. +// Unlike nlp.QuantLlama (Llama-only), this reaches the Qwen2 (qkv-bias) and Qwen3 (QK-norm) variants — +// biases and QK-norm gains stay f32 and apply after the Q8 matmul. Q8 is not bit-exact (cosine-validated). +func NewLlamaQ8CUDA(m *nlp.Llama) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + fusedGateUp: true, // cRec implements SwiGLUHalves; newDecoder column-fuses ffn_gate|ffn_up + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + quantizeF32: func(w *tensor.Tensor) (qweight, error) { + return cuda.NewResidentBQ8(w) + }, + }) +} + +// cudaQ4KResident quantizes a host f32 weight [K,N] to a resident Q4_K (ggml 4-bit k-quant, 256-element +// affine super-blocks) device weight — the aggressive-quant option. It transposes [K,N]→[N,K] (Q4_K blocks +// are per output row), quantizes with the tested gguf.Quantize(Q4_K), and uploads via NewResidentBQ4KFromBlocks. +// Q4_K is 2× smaller than Q8 (a 13B model fits the 12GB 3060 at Q4_K but not Q8) and matches ggml Q4_K_M +// accuracy — but decode-only (no int-tensor-core prefill), so prefill is slower than Q8. K must be a mult of 256. +func cudaQ4KResident(w *tensor.Tensor) (qweight, error) { + kk, nn := w.Shape()[0], w.Shape()[1] + if kk%256 != 0 { + return nil, fmt.Errorf("llamagpu: Q4_K needs K%%256==0, got K=%d (weight %dx%d)", kk, kk, nn) + } + bf := w.Cast(tensor.F32).Contiguous().Storage().F32() // [K,N] row-major: bf[r*nn+col] + wt := make([]float32, nn*kk) // [N,K] + for col := 0; col < nn; col++ { + base := col * kk + //perfscan:ignore PS4004 one-time Q4_K weight repack/quantize at model load + for r := 0; r < kk; r++ { + //perfscan:ignore PS6011 load-time resident-block construction; cold + wt[base+r] = bf[r*nn+col] + } + } + t := tensor.New(tensor.F32, tensor.Shape{nn, kk}) + copy(t.Storage().F32(), wt) + blocks, err := gguf.Quantize(t, gguf.Q4_K) + if err != nil { + return nil, err + } + return cuda.NewResidentBQ4KFromBlocks(blocks, kk, nn) +} + +// NewLlamaQ4KCUDA is NewCUDA with every projection quantized to resident Q4_K (4-bit k-quant) — 2× smaller +// than Q8, matching ggml Q4_K_M accuracy, with a direct f32 path (no ggml file needed). It lets the 12GB +// 3060 run models that don't fit at Q8. Decode is a Q4_K GEMV (weight-bandwidth win over Q8); prefill has no +// int-tensor-core GEMM so is slower than Q8's MMQ — prefer Q8 when prefill throughput dominates. Reaches the +// Qwen2/Qwen3 variants nlp.QuantLlama can't. K%256==0 required (typical transformer dims satisfy it). cuda-only. +func NewLlamaQ4KCUDA(m *nlp.Llama) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + fusedGateUp: true, // cRec implements SwiGLUHalves; newDecoder column-fuses ffn_gate|ffn_up + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + quantizeF32: cudaQ4KResident, + }) +} + +// NewQwen2Q4KCUDA / NewQwen3Q4KCUDA / NewPhi3Q4KCUDA are NewLlamaQ4KCUDA typed entry points for the popular +// Llama variants (Qwen2/2.5 qkv-bias, Qwen3 QK-norm, Phi-3) — their 4-bit path so a bigger model fits the +// 12GB 3060. Load with nlp.LlamaFromHF / Phi3FromHF; biases/QK-norm stay f32. +func NewQwen2Q4KCUDA(m *nlp.Llama) (*Decoder, error) { return NewLlamaQ4KCUDA(m) } + +// NewQwen3Q4KCUDA is NewQwen2Q4KCUDA's Qwen3 sibling (see above): NewLlamaQ4KCUDA with a Qwen3-typed +// signature, so it picks up Qwen3's per-head QK-norm automatically from the checkpoint. +func NewQwen3Q4KCUDA(m *nlp.Llama) (*Decoder, error) { return NewLlamaQ4KCUDA(m) } + +// NewPhi3Q4KCUDA is NewQwen2Q4KCUDA's Phi-3 sibling (see above): NewLlamaQ4KCUDA with a Phi-3-typed +// signature. Load with nlp.Phi3FromHF. +func NewPhi3Q4KCUDA(m *nlp.Llama) (*Decoder, error) { return NewLlamaQ4KCUDA(m) } + +// NewQwen2Q8CUDA / NewQwen3Q8CUDA / NewPhi3Q8CUDA are NewLlamaQ8CUDA typed entry points for the Llama +// variants that nlp.QuantLlama cannot represent (Qwen2 qkv-bias, Qwen3 QK-norm) or that only had an f32 +// GPU path (Phi-3). They give those models their first quantized decode. Load with nlp.LlamaFromHF / +// Phi3FromHF; biases and QK-norm gains are picked up and kept f32. +func NewQwen2Q8CUDA(m *nlp.Llama) (*Decoder, error) { return NewLlamaQ8CUDA(m) } + +// NewQwen3Q8CUDA is NewQwen2Q8CUDA's Qwen3 sibling (see above): NewLlamaQ8CUDA with a Qwen3-typed +// signature, so it picks up Qwen3's per-head QK-norm automatically from the checkpoint. +func NewQwen3Q8CUDA(m *nlp.Llama) (*Decoder, error) { return NewLlamaQ8CUDA(m) } + +// NewPhi3Q8CUDA is NewQwen2Q8CUDA's Phi-3 sibling (see above): NewLlamaQ8CUDA with a Phi-3-typed +// signature — Phi-3's first quantized GPU decode path. Load with nlp.Phi3FromHF. +func NewPhi3Q8CUDA(m *nlp.Llama) (*Decoder, error) { return NewLlamaQ8CUDA(m) } + +// cudaQ8Ops is the cuda backendOps with resident-Q8_0 quantization enabled (ops.quantizeF32 set) — +// shared by the dense-transformer NewXQ8CUDA entry points below. Each routes its arch's f32 checkpoint +// through the same weight-bandwidth lever (~2-3× faster decode): mkLin/mkFused quantize every projection +// and fused QKV; biases, QK-norm, ALiBi and RoPE stay f32. Q8 is not bit-exact (cosine-validated). These +// reach arches with NO prior quant path (nlp.QuantLlama is Llama-only). +func cudaQ8Ops() backendOps { + return backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + quantizeF32: func(w *tensor.Tensor) (qweight, error) { return cuda.NewResidentBQ8(w) }, + } +} + +// NewGemma2Q8CUDA is NewGemma2CUDA with every projection (fused QKV, o_proj, gate/up/down) and the +// tied lm_head quantized to resident Q8_0 — Gemma2's vocab is huge (256k), so newGemma2Decoder routes +// the tied head through the same quantize-aware builder rather than leaving it f32. The attention-logit +// soft-cap (MHACap), √dim embed scaling and (1+w) RMSNorm gains stay f32. cuda-only. +func NewGemma2Q8CUDA(m *nlp.Gemma2) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newGemma2Decoder(m, cudaQ8Ops()) +} + +// NewGemmaQ8CUDA is NewGemmaCUDA (Gemma v1) with all projections + the tied lm_head quantized to resident +// Q8_0. Gemma's √dim embed scaling and GeGLU are unchanged; the lm_head goes Q8 too (huge 256k vocab). +func NewGemmaQ8CUDA(m *nlp.Gemma) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + o := cudaQ8Ops() + o.fusedGateUp = true // cRec implements GeGLUHalves; newGemmaDecoder column-fuses ffn_gate|ffn_up + return newGemmaDecoder(m, o) +} + +// NewGraniteQ8CUDA is the Q8 entry point for IBM Granite (dense) — Granite maps to nlp.Llama (Llama core + +// embMult/attn-scale/residual-mult/logit-scale config scalars, all folded into the upload), so this is +// NewLlamaQ8CUDA with a Granite-typed signature. The scaled Wo/Wdown are quantized post-scale. +func NewGraniteQ8CUDA(m *nlp.Llama) (*Decoder, error) { return NewLlamaQ8CUDA(m) } + +// NewCohereQ8CUDA is NewCohereCUDA with every projection (fused QKV, o_proj, gate/up/down) quantized to +// resident Q8_0 — Command-R's weight-bandwidth win. The logit_scale-folded tied lm_head is quantized +// too (via mkLinS); the parallel-residual weight-only mean-centered LayerNorm gains stay f32. cuda-only. +func NewCohereQ8CUDA(m *nlp.Cohere) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + o := cudaQ8Ops() + o.fusedGateUp = true // cRec implements SwiGLUHalves; builder column-fuses ffn_gate|ffn_up + return newCohereDecoder(m, o) +} + +// NewNemotronQ8CUDA is NewNemotronCUDA with every projection (fused QKV, o_proj, the 2-layer relu² +// MLP's up/down) and the lm_head quantized to resident Q8_0. The LayerNorm1P gain/bias (γ=w+1, β +// folded in by NemotronFromHF) and the squared-ReLU activation stay f32. cuda-only. +func NewNemotronQ8CUDA(m *nlp.Nemotron) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newNemotronDecoder(m, cudaQ8Ops()) +} + +// NewOLMo2Q8CUDA is NewOLMo2CUDA with every projection (fused QKV, o_proj, gate/up/down) and the +// untied lm_head quantized to resident Q8_0. The post-norm placement and the full-width q/k RMSNorm +// gains (applied before RoPE) stay f32. cuda-only. +func NewOLMo2Q8CUDA(m *nlp.OLMo2) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + o := cudaQ8Ops() + o.fusedGateUp = true // cRec implements SwiGLUHalves; builder column-fuses ffn_gate|ffn_up + return newOLMo2Decoder(m, o) +} + +// NewFalconQ8CUDA is NewFalconCUDA with every projection (fused QKV, dense, the 2-layer GELU MLP's +// up/down) and the lm_head quantized to resident Q8_0. The parallel-residual LayerNorm-with-bias gains +// stay f32. cuda-only. +func NewFalconQ8CUDA(m *nlp.Falcon) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newFalconDecoder(m, cudaQ8Ops()) +} + +// NewStableLMQ8CUDA is NewStableLMCUDA with every projection (fused QKV, o_proj, gate/up/down) and the +// lm_head quantized to resident Q8_0. The LayerNorm-with-bias gains and the partial-rotary split stay +// f32. cuda-only. +func NewStableLMQ8CUDA(m *nlp.StableLM) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + o := cudaQ8Ops() + o.fusedGateUp = true // cRec implements SwiGLUHalves; builder column-fuses ffn_gate|ffn_up + return newStableLMDecoder(m, o) +} + +// NewStarCoder2Q8CUDA is NewStarCoder2CUDA with every projection (fused QKV, o_proj, the 2-layer GELU +// MLP's fc/proj) and the untied lm_head quantized to resident Q8_0. Every projection bias and the +// LayerNorm-with-bias gains stay f32. cuda-only. +func NewStarCoder2Q8CUDA(m *nlp.StarCoder2) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newStarCoder2Decoder(m, cudaQ8Ops()) +} + +// NewMPTQ8CUDA is NewMPTCUDA with every projection and the tied lm_head quantized to resident Q8_0. +// The per-head ALiBi slopes and the weight-only LayerNorm gains stay f32. cuda-only. +func NewMPTQ8CUDA(m *nlp.MPT) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newMPTDecoder(m, cudaQ8Ops()) +} + +// NewGPTNeoXQ8CUDA is NewGPTNeoXCUDA with every projection (fused QKV, o_proj, the 2-layer GELU MLP's +// up/down) and the untied, unbiased lm_head quantized to resident Q8_0. Every projection bias, the +// two-norm LayerNorm gains and the (partial or full) rotary split stay f32. cuda-only. +func NewGPTNeoXQ8CUDA(m *nlp.GPTNeoX) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newGPTNeoXDecoder(m, cudaQ8Ops()) +} + +// NewPhiQ8CUDA is NewPhiCUDA with every projection (fused QKV, dense, the 2-layer GELU MLP's fc1/fc2) +// and the untied lm_head quantized to resident Q8_0. Every projection bias (including the lm_head's), +// the final LayerNorm and the partial-rotary split stay f32. cuda-only. +func NewPhiQ8CUDA(m *nlp.Phi) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newPhiDecoder(m, cudaQ8Ops()) +} + +// The MoE Q8 entry points quantize the (bandwidth-dominant) expert gate/up/down matrices and attention +// projections to resident Q8_0 while the top-k ROUTER stays f32 (d.f32Lin) — Q8 rounding in the router +// logits could flip expert selection. Sparse-MoE decode streams many expert matrices per token, so this +// is the largest Q8 decode win. Qwen3-MoE loads as an nlp.Mixtral, so NewQwen3MoEQ8CUDA == NewMixtralQ8CUDA. +// cudaQ4KOps is cudaQ8Ops with resident-Q4_K quantization (0.5625 B/weight vs Q8 1.0625 — ~2x less +// VRAM). Every projection with K%256==0 becomes a resident Q4_K weight; MoE experts (ResidentBQ4K) then +// get the gated sparse decode (cu_qmatmul_q4k_moe) automatically via cRec.QMatMulResidentMoE. +func cudaQ4KOps() backendOps { + o := cudaQ8Ops() + o.quantizeF32 = cudaQ4KResident + return o +} + +// NewMixtralQ4KCUDA uploads an nlp.Mixtral with Q4_K-quantized weights (experts included) — the Q4_K_M +// MoE path: ~2x less VRAM than NewMixtralQ8CUDA plus the ~E/K gated expert decode. Requires every +// projection K (=Dim; per-expert Dim/Hidden) %256==0. Qwen3-MoE loads as an nlp.Mixtral, so it applies too. +func NewMixtralQ4KCUDA(m *nlp.Mixtral) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + o := cudaQ4KOps() + o.fusedGateUp = true // cRec implements SwiGLUHalves; MoE experts column-fuse gate|up + return newMixtralDecoder(m, o) +} + +// NewMixtralQ8CUDA uploads an nlp.Mixtral with Q8_0-quantized weights (experts included) — the +// higher-fidelity counterpart to [NewMixtralQ4KCUDA]: ~2x the VRAM for weights that dequantize to +// within Q8_0's much tighter error bound, over the same ~E/K gated expert decode. Prefer it when a +// Q4_K MoE loses accuracy on the target task and the model still fits. Returns an error when no +// CUDA GPU is available. +func NewMixtralQ8CUDA(m *nlp.Mixtral) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + o := cudaQ8Ops() + o.fusedGateUp = true // cRec implements SwiGLUHalves; MoE experts column-fuse gate|up + return newMixtralDecoder(m, o) +} + +// NewQwen3MoEQ8CUDA is NewMixtralQ8CUDA's Qwen3-MoE entry point (see the MoE Q8 note above) — +// Qwen3-MoE loads as an nlp.Mixtral, so this is exactly NewMixtralQ8CUDA. Load with nlp.Qwen3MoeFromHF. +func NewQwen3MoEQ8CUDA(m *nlp.Mixtral) (*Decoder, error) { return NewMixtralQ8CUDA(m) } + +// NewQwen2MoEQ8CUDA is the Q8 entry point for nlp.Qwen2MoE (see the MoE Q8 note above): the routed and +// shared-expert matrices plus the Qwen2 attention (with its q/k/v projection biases) go resident Q8_0; +// the shared-expert sigmoid gate is quantized too. The top-k router stays f32. cuda-only. +func NewQwen2MoEQ8CUDA(m *nlp.Qwen2MoE) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + o := cudaQ8Ops() + o.fusedGateUp = true + return newQwen2MoEDecoder(m, o) +} + +// NewGraniteMoEQ8CUDA is the Q8 entry point for nlp.GraniteMoE (see the MoE Q8 note above): the +// attention projections and the routed expert gate/up/down matrices go resident Q8_0, with +// ResidualMult/LogitsScale folded pre-quantization into Wo/expert-Wdown/lm_head (as NewGraniteQ8CUDA +// does for dense Granite). The top-k router stays f32 (routing is selection-sensitive). cuda-only. +func NewGraniteMoEQ8CUDA(m *nlp.GraniteMoE) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + o := cudaQ8Ops() + o.fusedGateUp = true + return newGraniteMoEDecoder(m, o) +} + +// NewOLMoEQ8CUDA is the Q8 entry point for nlp.OLMoE (see the MoE Q8 note above): the attention +// projections and the routed expert matrices go resident Q8_0; the full-width q/k RMSNorm gains and +// the top-k router stay f32. cuda-only. +func NewOLMoEQ8CUDA(m *nlp.OLMoE) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + o := cudaQ8Ops() + o.fusedGateUp = true + return newOLMoEDecoder(m, o) +} + +// NewDeepSeekV2Q8CUDA quantizes the flagship MLA + DeepSeekMoE decoder — the last Q8 gap. The low-rank +// MLA projections (q_a/q_b/kv_a/kv_b, o_proj), the routed + shared expert matrices all go resident Q8_0; +// the latent RMSNorms, decoupled RoPE and the top-k router stay f32 (the router carved out, as for the +// other MoE arches). Completes Q8 decode across the entire f32 transformer catalogue. +func NewDeepSeekV2Q8CUDA(m *nlp.DeepSeekV2) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + o := cudaQ8Ops() + o.fusedGateUp = true + return newDeepSeekV2Decoder(m, o) +} + +// NewQwen2CUDA uploads a Qwen2 / Qwen2.5 model onto the batched Decoder core. Qwen2 shares +// nlp.Llama (SwiGLU MLP, RMSNorm, GQA, full rope) and departs only in carrying q/k/v projection +// biases (o_proj has none) — the newDecoder core adds them via qkvBias when b.Bq/Bk/Bv are set, so +// this is exactly NewCUDA with a Qwen2-typed entry point. Load with nlp.LlamaFromHF; the biases are +// picked up automatically from the checkpoint. (Qwen3's per-head QK-norm is a separate follow-up.) +func NewQwen2CUDA(m *nlp.Llama) (*Decoder, error) { return NewCUDA(m) } + +// NewQwen3CUDA uploads a Qwen3 model onto the batched Decoder core. Qwen3 shares nlp.Llama and adds +// per-head RMSNorm on Q and K before RoPE (b.QNorm/b.KNorm), which newDecoder wires into recordQKNorm +// — it drops Qwen2's q/k/v projection bias, so this is again NewCUDA with a Qwen3-typed entry point. +// Load with nlp.LlamaFromHF; the QK-norm gains are picked up automatically from the checkpoint. +func NewQwen3CUDA(m *nlp.Llama) (*Decoder, error) { return NewCUDA(m) } + +// NewPhi3CUDA uploads a Phi-3 model onto the batched Decoder core. Phi-3 is structurally a plain +// Llama (RMSNorm, SwiGLU, GQA, full rope, no biases) — nlp.Phi3FromHF just unpacks its row-packed +// qkv_proj / gate_up_proj into the standard projections and returns an *nlp.Llama — so this is +// NewCUDA with a Phi-3-typed entry point. Distinct from NewPhiCUDA (the older Phi-1/1.5/2 family, +// which is parallel-residual with LayerNorm and partial rotary). +func NewPhi3CUDA(m *nlp.Llama) (*Decoder, error) { return NewCUDA(m) } + +// NewNemotronCUDA uploads an nlp.Nemotron onto the batched Decoder core: sequential two-norm +// residual with LayerNorm1P (mean-centered LayerNorm, γ=w+1/β folded at load), a squared-ReLU +// 2-layer MLP (down(relu²(up(x)))) and partial rotary. The relu² activation is a new cuda unary +// (cu_relu2_f32); everything else is reused plumbing. cuda-only. +func NewNemotronCUDA(m *nlp.Nemotron) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newNemotronDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewFalconCUDA uploads an nlp.Falcon (falcon-7b class) onto the batched Decoder core: single-norm +// parallel residual, LayerNorm-with-bias, a 2-layer GELU MLP, multi-query attention (one KV head) +// and full rope — every departure a reused generalization. cuda-only. +func NewFalconCUDA(m *nlp.Falcon) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newFalconDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewOLMo2CUDA uploads an nlp.OLMo2 (Allen AI) onto the batched Decoder core: post-norm blocks +// (each sublayer reads the raw residual, its output normed before the add), a full-width RMSNorm on +// the q/k projections before RoPE, SwiGLU, GQA, full rope and an untied lm_head. cuda-only. +func NewOLMo2CUDA(m *nlp.OLMo2) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newOLMo2Decoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewGraniteMoECUDA uploads an nlp.GraniteMoE onto the batched Decoder core: the MoE sibling of +// dense Granite — a plain Llama attention core + sparse Mixture-of-Experts FFN, wrapped in the four +// Granite config scalars (embedding/attention/residual multipliers folded into the upload, logits +// scaling into the lm_head). cuda-only. +func NewGraniteMoECUDA(m *nlp.GraniteMoE) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newGraniteMoEDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewQwen2MoECUDA uploads an nlp.Qwen2MoE onto the batched Decoder core: a routed sparse MoE PLUS a +// shared expert (a SwiGLU run on every token, sigmoid-gated), with Qwen2 q/k/v projection biases on +// the attention. cuda-only. +func NewQwen2MoECUDA(m *nlp.Qwen2MoE) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newQwen2MoEDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewOLMoECUDA uploads an nlp.OLMoE (Allen AI sparse-MoE) onto the batched Decoder core: pre-norm +// Llama attention with FULL-WIDTH q/k RMSNorm and a sparse Mixture-of-Experts FFN, untied lm_head. +// cuda-only. +func NewOLMoECUDA(m *nlp.OLMoE) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newOLMoEDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewDeepSeekV2CUDA uploads a DENSE (MLA-only) nlp.DeepSeekV2 onto the batched Decoder core: +// Multi-head Latent Attention (low-rank latent KV compression, decoupled RoPE, rectangular +// query/key vs value head dims) + SwiGLU FFN. The hardest attention in the catalogue, and the first +// non-fused-QKV decoder. cuda-only (the rectangular MHARect is cuda-only). The sparse-MoE variant is +// a follow-up (dense-only for now). +func NewDeepSeekV2CUDA(m *nlp.DeepSeekV2) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newDeepSeekV2Decoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewMambaCUDA uploads an nlp.Mamba onto the batched Decoder core — the FIRST non-transformer GPU +// decoder. Each layer is a selective-state-space mixer (no attention, no KV cache): decode is a +// linear-time recurrence recorded from the per-timestep conv/SSM decode kernels (cu_conv1d_step / +// cu_ssm_step), carrying per-block conv/SSM state across Step calls. cuda-only. +func NewMambaCUDA(m *nlp.Mamba) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newMambaDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewMambaQ8CUDA is NewMambaCUDA with the projection weights quantized to resident Q8_0. Mamba decode +// is weight-bandwidth-bound (~230 MB of f32 weights streamed per token), so the Q8_0 GEMVs cut the +// dominant cost ~4× at the price of Q8 rounding (not bit-exact — validated to a tolerance, not parity). +// The tiny per-channel state (A, dt_bias, conv, dskip, norms) stays f32. cuda-only. +func NewMambaQ8CUDA(m *nlp.Mamba) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newMambaDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + quantizeF32: func(w *tensor.Tensor) (qweight, error) { + return cuda.NewResidentBQ8(w) + }, + }) +} + +// NewJambaCUDA uploads an nlp.Jamba (hybrid Mamba + NoPE-attention + MoE/dense) onto the batched +// Decoder core. Each layer is a Mamba selective-scan mixer OR grouped-query causal attention (no rotary +// — the Mamba layers carry position), followed by a sparse-MoE or dense SwiGLU FFN. Attention layers +// keep a growing KV cache; Mamba layers carry conv/SSM state across Step, so the whole model decodes +// sequentially (rows==1 Step; StepN loops Step). cuda-only. +func NewJambaCUDA(m *nlp.Jamba) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newJambaDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewJambaQ8CUDA is NewJambaCUDA with the projection weights quantized to resident Q8_0 — the hybrid's +// Mamba mixers, NoPE-attention q/k/v/o, MoE expert gate/up/down, dense SwiGLU and lm_head all go Q8 via +// the shared mkLin, while the MoE top-k router stays f32 (d.f32Lin) so routing is exact. Jamba decodes +// rows==1 (Step), so this is a decode-bandwidth win (the GEMV path); StepN loops Step. cuda-only. +func NewJambaQ8CUDA(m *nlp.Jamba) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newJambaDecoder(m, cudaQ8Ops()) +} + +// NewMamba2CUDA uploads an nlp.Mamba2 onto the batched Decoder core — the state-space-duality sibling +// of Mamba. Each layer is an SSD mixer (scalar per-head decay, B/C shared across a group, gated +// RMSNorm); decode is a linear-time recurrence recorded from cu_ssd_step + the shared conv1d/softplus +// primitives, carrying per-block conv/SSD state across Step calls. cuda-only. +func NewMamba2CUDA(m *nlp.Mamba2) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newMamba2Decoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewMamba2Q8CUDA is NewMamba2CUDA with the in/out projections and LM head quantized to resident Q8_0 — +// the same weight-bandwidth lever as NewMambaQ8CUDA. Mamba-2's InProj/OutProj are stored [out,in] (torch +// layout), so newMamba2Decoder transposes each to [in,out] before quantizing (NewResidentBQ8 materializes +// the transposed view). The SSD recurrence (per-head decay, conv, gated RMSNorm) stays f32. cuda-only. +func NewMamba2Q8CUDA(m *nlp.Mamba2) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newMamba2Decoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + quantizeF32: func(w *tensor.Tensor) (qweight, error) { + return cuda.NewResidentBQ8(w) + }, + }) +} + +// NewRWKVCUDA uploads an nlp.RWKV onto the batched Decoder core — GoAI's third recurrent family. Each +// layer is a WKV time-mix (recorded from cu_wkv_step) + a gated squared-ReLU channel-mix; decode is +// O(1) recurrent with no attention and no KV cache, carrying per-block token-shift + WKV state across +// Step calls. cuda-only. +func NewRWKVCUDA(m *nlp.RWKV) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newRWKVDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewRWKVQ8CUDA is NewRWKVCUDA with the time-mix (Wr/Wk/Wv/Wo) and channel-mix (CWr/CWk/CWv) GEMVs plus +// the LM head quantized to resident Q8_0 — the same weight-bandwidth lever as NewMambaQ8CUDA. RWKV decode +// streams these projection weights per token; Q8_0 cuts that ~4× at the price of Q8 rounding (validated to +// a cosine tolerance, not parity). The tiny per-channel mix params (μ/decay/bonus) and LayerNorms stay f32. +func NewRWKVQ8CUDA(m *nlp.RWKV) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newRWKVDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + quantizeF32: func(w *tensor.Tensor) (qweight, error) { + return cuda.NewResidentBQ8(w) + }, + }) +} + +// NewBertCUDA uploads an nlp.Bert bidirectional encoder onto the GPU — the first non-decoder GPU +// model in llamagpu. It runs one bidirectional forward (post-LN, learned absolute positions, no +// causal mask, no KV cache) returning the [seq, dim] hidden states, reusing the existing recorder +// ops (MHA with causal=0, LayerNorm, GELU MLP) with no new kernel. cuda-only. +func NewBertCUDA(m *nlp.Bert) (*GPUBert, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newBertEncoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewBertQ8CUDA is NewBertCUDA with the attention (q/k/v/o) and FFN (w1/w2) projections quantized to +// resident Q8_0. BERT encodes the whole sequence at once (M=seq>1), so each Q8 projection runs on the +// int8 tensor-core MMQ GEMM (QMatMulResident's m>1 path) — faster than the f32 cuBLAS GEMM and 4× less +// weight memory. Biases and LayerNorm stay f32. Works for RoBERTa/DistilBERT too (all *nlp.Bert). cuda-only. +func NewBertQ8CUDA(m *nlp.Bert) (*GPUBert, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newBertEncoder(m, cudaQ8Ops()) +} + +// NewT5CUDA uploads an nlp.T5 bidirectional encoder onto the GPU. T5 is the second non-decoder GPU +// model: PRE-LN residuals, RMSNorm (T5LayerNorm), a learned per-head RELATIVE-position bias added to +// the scores (via cu_attn_softmax_bias / MHABias), NO 1/√d scaling, NO absolute/rotary position, a +// gated-GELU (v1.1) or ReLU (v1.0) FFN and no projection biases. Returns [seq, dim] hidden states. +// cuda-only (MHABias is cuda-only). +func NewT5CUDA(m *nlp.T5) (*GPUT5, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newT5Encoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewT5Q8CUDA is NewT5CUDA with the attention (q/k/v/o) and FFN (wi0/wi1/wOut) projections quantized to +// resident Q8_0. T5 encodes the whole sequence at once (M=seq>1), so each Q8 projection runs on the int8 +// tensor-core MMQ GEMM (QMatMulResident's m>1 path) — faster than the f32 cuBLAS GEMM and 4× less weight +// memory. RMSNorm and the relative-position bias stay f32. Covers T5 v1.0 (ReLU) and v1.1 (GeGLU). cuda-only. +func NewT5Q8CUDA(m *nlp.T5) (*GPUT5, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newT5Encoder(m, cudaQ8Ops()) +} + +// NewT5DecoderCUDA uploads an nlp.T5Decoder onto the GPU — the seq2seq (encoder-decoder) decoder, the +// first encoder-decoder GPU model. Each block runs causal self-attention with the T5 relpos bias +// (MHABias over a growing KV cache), cross-attention over the encoder output (plain MHA), and a +// gated-GELU/ReLU FFN, all PRE-LN/RMSNorm/unscaled. Pair with NewT5CUDA (the encoder). cuda-only. +func NewT5DecoderCUDA(m *nlp.T5Decoder) (*GPUT5Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newT5Decoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewT5DecoderQ8CUDA is NewT5DecoderCUDA with the self/cross-attention (q/k/v/o), FFN (wi0/wi1/wOut) and +// lm_head projections quantized to resident Q8_0 — the seq2seq generation decode-bandwidth win. Autoregressive +// decode is M=1 (Q8 GEMV); the cross K/V projections run once over the encoder output (M=eseq>1 → int8 +// tensor-core MMQ). RMSNorm and the relpos bias stay f32. Pair with NewT5Q8CUDA (the Q8 encoder). cuda-only. +func NewT5DecoderQ8CUDA(m *nlp.T5Decoder) (*GPUT5Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newT5Decoder(m, cudaQ8Ops()) +} + +// NewMixtralCUDA uploads an nlp.Mixtral (sparse Mixture-of-Experts) onto the batched Decoder core. +// Attention is the plain Llama core (+ optional Qwen3-MoE QK-norm); the FFN is a sparse MoE routed +// per token — the routing weights are computed on-device (cu_moe_gate) so the whole step stays a +// pre-recorded batched command buffer. Dense expert eval (correct); sparse top-k gather is a +// follow-up. cuda-only. +func NewMixtralCUDA(m *nlp.Mixtral) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newMixtralDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewQwen3MoECUDA uploads a Qwen3-MoE model onto the batched Decoder core. Qwen3-MoE loads as an +// nlp.Mixtral whose blocks additionally carry per-head q/k RMSNorm gains (the same optional QK-norm +// newMixtralDecoder already wires), so this is exactly NewMixtralCUDA with a Qwen3-MoE entry point. +// Load with nlp.Qwen3MoeFromHF. +func NewQwen3MoECUDA(m *nlp.Mixtral) (*Decoder, error) { return NewMixtralCUDA(m) } + +// NewMPTCUDA uploads an nlp.MPT (MosaicML) onto the batched Decoder core: the first ALiBi decoder — +// position enters solely through a per-head linear bias on the attention scores (no RoPE), with +// weight-only LayerNorm, standard MHA, a bias-free 2-layer GELU MLP and a tied lm_head. cuda-only +// (the ALiBi attention kernel is cuda-only). +func NewMPTCUDA(m *nlp.MPT) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newMPTDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewGemma2CUDA uploads an nlp.Gemma2 onto the batched Decoder core: Gemma's √dim embed / (1+w) +// RMSNorm / GeGLU / decoupled head_dim / tied lm_head, plus sandwich norms (pre + post per sublayer), +// an attention-logit soft-cap (via MHACap) and query_pre_attn_scalar. The final-logit soft-cap is +// monotonic (greedy-invariant) and omitted. cuda-only (the soft-cap kernel is cuda-only). +func NewGemma2CUDA(m *nlp.Gemma2) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newGemma2Decoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewGemmaCUDA uploads an nlp.Gemma (Gemma v1) onto the batched Decoder core: RMSNorm (with the +// (1+w) gain folded at load), RoPE, GQA, a √dim embedding normalizer (d.embMult), a GeGLU FFN and +// a tied lm_head — every departure a reused generalization or a load-time fold. cuda-only. +func NewGemmaCUDA(m *nlp.Gemma) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newGemmaDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewCohereCUDA uploads an nlp.Cohere (Command-R) onto the batched Decoder core: one-norm parallel +// residual, weight-only mean-centered LayerNorm, SwiGLU, GQA, full rope (its interleaved rotary is +// pre-permuted into the q/k weights by CohereFromHF) and a logit_scale-folded tied lm_head — every +// piece a reused generalization. cuda-only (parallel residual + LayerNorm are the Phi plumbing). +func NewCohereCUDA(m *nlp.Cohere) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newCohereDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewGraniteCUDA uploads an IBM Granite model onto the batched Decoder core. Granite is a plain +// Llama plus four learned-at-config scalars — embedding_multiplier, attention_multiplier, +// residual_multiplier and logits_scaling — which newDecoder folds into the upload (embedding gather +// scale, softmax scale override, Wo/Wdown pre-scale, and 1/logits_scaling into the lm_head), so this +// is NewCUDA with a Granite-typed entry point. Load with nlp.GraniteFromHF. +func NewGraniteCUDA(m *nlp.Llama) (*Decoder, error) { return NewCUDA(m) } + +// NewStableLMCUDA uploads an nlp.StableLM into CUDA device buffers and runs it through the same +// batched Decoder core as NewCUDA — but with LayerNorm-with-bias norms and PARTIAL rotary (the +// StableLM/Phi/StarCoder2-class departures from Llama). The first of the new-architecture GPU +// graph decoders enabled by the cu_rope_partial* kernels. cuda-only for now. +func NewStableLMCUDA(m *nlp.StableLM) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newStableLMDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewStarCoder2CUDA uploads an nlp.StarCoder2 onto the batched Decoder core: LayerNorm-with-bias, +// biased q/k/v/o projections, a biased 2-layer GELU MLP, full rope and GQA. The second new-arch GPU +// graph decoder, exercising every Decoder-core generalization except partial rotary. cuda-only. +func NewStarCoder2CUDA(m *nlp.StarCoder2) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newStarCoder2Decoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewPhiCUDA uploads an nlp.Phi onto the batched Decoder core: one-norm parallel residual, biased +// q/k/v/dense projections, a biased GELU MLP, partial rotary, a biased untied lm_head and a final +// LayerNorm — every Decoder-core generalization at once. The third new-arch GPU graph decoder. +func NewPhiCUDA(m *nlp.Phi) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newPhiDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewGPTNeoXCUDA uploads an nlp.GPTNeoX onto the batched Decoder core: TWO-norm parallel residual +// (input_layernorm → attn and post_attention_layernorm → mlp, both over the raw residual), +// LayerNorm-with-bias, biased q/k/v/dense projections, a biased GELU MLP, (partial or full) rotary +// and an untied unbiased lm_head. The fourth new-arch GPU graph decoder — the two-norm sibling of +// Phi's one-norm parallel residual. cuda-only. +func NewGPTNeoXCUDA(m *nlp.GPTNeoX) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newGPTNeoXDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// NewGPTCUDA uploads an nlp.GPT (GPT-2-style: learned positional embeddings, LayerNorm, GELU MLP) +// into CUDA device buffers for batched decoding — the NVIDIA variant of NewGPT/NewGPTVulkan (§T422). +// Closes the last backend gap for the GPT-2 decoder, which previously ran only on metal and vulkan. +func NewGPTCUDA(m *nlp.GPT) (*GPTDecoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newGPTDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + }) +} + +// cudaUploadQWeight makes a ggml quantized [Out,In] weight resident on the GPU. +// +// Q4_K sources stay NATIVE: the raw GGUF Q4_K blocks (row-major [Out][In/256] super-blocks — exactly +// what NewResidentBQ4KFromBlocks wants) upload as-is, keeping the format's 0.5625 B/weight (vs 1.0625 +// upcast to Q8) — ~2× less VRAM and less decode bandwidth, at IDENTICAL accuracy (Q4_K→f32→Q8 adds +// nothing a Q4_K weight didn't already lose). Q4_K's resident dispatch (QMatMulResidentQ4K) already +// serves the GEMV decode AND the tensor-core WMMA prefill, and the GGUF-block path is exactly the one +// TestCUDAQ4KMatMulParity validates against the exact-dequant reference. Any error falls through to Q8. +// +// Other source types (Q5_K, Q6_K, …) have no native CUDA resident dispatch yet, so they are +// dequantized to f32 and re-quantized to a uniform resident Q8_0 (a faithful, near-lossless path). +// NewResidentBQ8 wants [In,Out], so the dequantized [Out,In] is transposed. +func cudaUploadQWeight(weight []byte, qt uint32, n, k int) (qweight, error) { + if qt == 12 /* tQ4_K */ && k%256 == 0 { + if rw, err := cuda.NewResidentBQ4KFromBlocks(weight, k, n); err == nil { + return rw, nil + } + // fall through to the always-correct Q8 path on any error + } + if qt == 14 /* tQ6_K */ && k%256 == 0 { + // Q6_K (0.8203 B/w vs Q8 1.0625): keep native — its resident dispatch (QMatMulResidentQ6K) + // serves the GEMV decode AND the tensor-core WMMA prefill (#874/#880). Raw GGUF Q6_K blocks + // are row-major [Out][In/256] 210-byte super-blocks = NewResidentBQ6KFromBlocks's input. + if rw, err := cuda.NewResidentBQ6KFromBlocks(weight, k, n); err == nil { + return rw, nil + } + // fall through to Q8 on any error + } + if qt == 13 /* tQ5_K */ && k%256 == 0 { + // Q5_K (0.6875 B/w vs Q8 1.0625, ~1.5× less VRAM): native via QMatMulResidentQ5K (GEMV decode + // + tensor-core WMMA prefill). Raw GGUF Q5_K blocks are row-major [Out][In/256] 176-byte + // super-blocks = NewResidentBQ5KFromBlocks's input. + if rw, err := cuda.NewResidentBQ5KFromBlocks(weight, k, n); err == nil { + return rw, nil + } + // fall through to Q8 on any error + } + if qt == 11 /* tQ3_K */ && k%256 == 0 { + // Q3_K (0.43 B/w vs Q8 1.0625, ~2.5× less VRAM): native via QMatMulResidentQ3K (GEMV decode + + // tensor-core WMMA prefill). NewResidentBQ3KFromBlocks repacks the raw GGUF Q3_K 110-byte blocks + // (row-major [Out][In/256]) into the aligned meta/qs/hmask planes the kernels want. + if rw, err := cuda.NewResidentBQ3KFromBlocks(weight, k, n); err == nil { + return rw, nil + } + // fall through to Q8 on any error + } + if qt == 10 /* tQ2_K */ && k%256 == 0 { + // Q2_K (0.25 B/w vs Q8 1.0625, ~4× less VRAM): the aggressive-quant option — a Q2_K GGUF upcast + // to Q8 may not even fit the 12GB 3060, so keeping it native is enabling, not just faster. Native + // via QMatMulResidentQ2K (GEMV decode + tensor-core WMMA prefill). Raw GGUF Q2_K blocks are + // row-major [Out][In/256] 84-byte super-blocks = NewResidentBQ2KFromBlocks's input. + if rw, err := cuda.NewResidentBQ2KFromBlocks(weight, k, n); err == nil { + return rw, nil + } + // fall through to Q8 on any error + } + if qt == 8 /* tQ8_0 */ && k%32 == 0 { + // Q8_0 stays NATIVE: the raw GGUF Q8_0 blocks (row-major [Out][In/32], f16 scale + 32 int8) repack + // DIRECTLY into ResidentBQ8 (q int8 + f32 scale) — no Q8_0→f32→Q8 round-trip. Cuts peak host RAM + // (a 7B Q8 that OOMs the f32 dequant path fits) and load work, at higher fidelity. Mirrors the + // k-quant FromBlocks paths above. Any error falls through to the always-correct f32 path. + if rw, err := cuda.NewResidentBQ8FromBlocks(weight, k, n); err == nil { + return rw, nil + } + } + f32, err := gguf.QuantTensor{Data: weight, GGType: qt, Shape: tensor.Shape{n, k}}.Dequantize() + if err != nil { + return nil, fmt.Errorf("llamagpu: CUDA dequant qt=%d [%d,%d]: %w", qt, n, k, err) + } + tin, err := f32.Transpose(0, 1) // [Out,In] → [In,Out] + if err != nil { + return nil, fmt.Errorf("llamagpu: CUDA weight transpose: %w", err) + } + rw, err := cuda.NewResidentBQ8(tin) + if err != nil { + return nil, fmt.Errorf("llamagpu: CUDA Q8 upload: %w", err) + } + return rw, nil +} + +// NewQuantCUDA uploads a quantized Llama with every projection as a resident Q8 weight +// (source types re-quantized to Q8) — the NVIDIA variant of NewQuant/NewQuantVulkan. +// Same batched Decoder core, so Generate/Step and every sampler work unchanged. +func NewQuantCUDA(m *nlp.QuantLlama) (*Decoder, error) { + if !cuda.Available() { + return nil, fmt.Errorf("llamagpu: no CUDA GPU") + } + return newQuantDecoder(m, backendOps{ + name: string(backend.CUDA), + asyncEncode: false, + fusedGateUp: true, // cRec implements SwiGLUHalves; newQuantDecoder column-fuses ffn_gate|ffn_up + newBuffer: func(data []float32) (buffer, error) { + b, err := cuda.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return cBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := cuda.NewRecorder() + if err != nil { + return nil, err + } + return cRec{r}, nil + }, + uploadQWeight: cudaUploadQWeight, + }) +} diff --git a/checks/testdata/ps6136-owner/export-api/main.go b/checks/testdata/ps6136-owner/export-api/main.go new file mode 100644 index 00000000..7dd20db2 --- /dev/null +++ b/checks/testdata/ps6136-owner/export-api/main.go @@ -0,0 +1,74 @@ +// Command export-api prints type-only model scaffolding from the pinned owner +// checkout. It is not backend execution or detector coverage evidence. +package main + +import ( + "fmt" + "go/types" + "os" + "sort" + + "golang.org/x/tools/go/packages" +) + +func main() { + pkgs, err := packages.Load(&packages.Config{Dir: os.Args[1], Mode: packages.NeedName | packages.NeedTypes | packages.NeedImports | packages.NeedDeps}, "github.com/jxsl13/goai/nlp") + if err != nil || packages.PrintErrors(pkgs) != 0 || len(pkgs) != 1 { + panic("cannot load pinned model metadata") + } + pkg := pkgs[0].Types + var selected *types.Package + selected = pkg + if len(os.Args) > 2 && os.Args[2] == "nn" { + for _, imported := range pkg.Imports() { + if imported.Path() == "github.com/jxsl13/goai/nn" { + selected = imported + } + } + } + seen := map[string]*types.Named{} + var visit func(types.Type) + visit = func(typ types.Type) { + switch value := typ.(type) { + case *types.Named: + if value.Obj().Pkg() != selected || seen[value.Obj().Name()] != nil { + return + } + seen[value.Obj().Name()] = value + visit(value.Underlying()) + case *types.Struct: + for index := 0; index < value.NumFields(); index++ { + visit(value.Field(index).Type()) + } + case *types.Pointer: + visit(value.Elem()) + case *types.Slice: + visit(value.Elem()) + case *types.Array: + visit(value.Elem()) + case *types.Map: + visit(value.Key()) + visit(value.Elem()) + } + } + namesToVisit := []string{"GPT", "Llama", "StableLM", "Cohere", "Nemotron", "Gemma", "DeepSeekV2", "Mamba", "Jamba", "Mamba2", "RWKV", "Falcon", "OLMo2", "Gemma2", "MPT", "Mixtral", "Qwen2MoE", "GraniteMoE", "OLMoE", "StarCoder2", "Phi", "GPTNeoX", "QuantLlama"} + if selected != pkg { + namesToVisit = []string{"RMSNorm", "LayerNorm", "SwiGLU", "GLU", "DeepSeekMoE", "LoRALinear", "MambaBlock", "SparseMoE", "QuantLinear", "QuantSwiGLU", "RWKVBlock"} + } + for _, name := range namesToVisit { + visit(selected.Scope().Lookup(name).Type()) + } + var names []string + for name := range seen { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + fmt.Printf("type %s %s\n", name, types.TypeString(seen[name].Underlying(), func(other *types.Package) string { + if other == selected { + return "" + } + return other.Name() + })) + } +} diff --git a/checks/testdata/ps6136-owner/llamagpu.go.txt b/checks/testdata/ps6136-owner/llamagpu.go.txt new file mode 100644 index 00000000..f577d550 --- /dev/null +++ b/checks/testdata/ps6136-owner/llamagpu.go.txt @@ -0,0 +1,450 @@ +//go:build darwin && cgo + +// Further reading: Pope et al. 2022, "Efficiently Scaling Transformer Inference", for the batched-decode systems view; Dao et al. 2022, "FlashAttention", for the memory-aware attention lineage behind the cooperative kernels. +// +// In plain terms: the fast lane for text generation — whole decode steps run as single pre-recorded GPU programs over weights that stay resident on the device, instead of shuttling tensors back and forth per operation. +package llamagpu + +import ( + "fmt" + "os" + + "github.com/jxsl13/goai/backend" + "github.com/jxsl13/goai/backend/metal" + "github.com/jxsl13/goai/nlp" +) + +// metal adapter (§T409): thin assertions from the backend-agnostic buffer/recorder interfaces to +// the concrete metal.DeviceBuffer/metal.Recorder (whose API vulkan mirrors method-for-method). + +type mBuf struct{ *metal.DeviceBuffer } + +type mRec struct{ r *metal.Recorder } + +func (m mRec) Barrier() error { return m.r.Barrier() } + +// mProfileRec is used only by ProfileMetalStep. It captures the completed recorder before the +// normal Step lifecycle frees it; embedding mRec keeps the production adapter and operation set +// identical. +type mProfileRec struct { + mRec + profile *metal.RecorderProfile +} + +func (m mRec) RMSNorm(x, g, o buffer, rows, dim int, eps float32) error { + return m.r.RMSNorm(mb(x), mb(g), mb(o), rows, dim, eps) +} +func (m mRec) LayerNorm(x, g, bb, o buffer, rows, dim int, eps float32) error { + return m.r.LayerNorm(mb(x), mb(g), mb(bb), mb(o), rows, dim, eps) +} +func (m mRec) AddBias(x, bb, o buffer, rows, n int) error { + return m.r.AddBias(mb(x), mb(bb), mb(o), rows, n) +} +func (m mRec) BiasGELU(x, bb, o buffer, rows, n int) error { + return m.r.BiasGELU(mb(x), mb(bb), mb(o), rows, n) +} +func (m mRec) MatMul(a, b, c buffer, mm, k, n int) error { + return m.r.MatMul(mb(a), mb(b), mb(c), mm, k, n) +} +func (m mRec) MatMulAcc(a, b, c buffer, mm, k, n int) error { + return m.r.MatMulAcc(mb(a), mb(b), mb(c), mm, k, n) +} +func (m mRec) F32QKVBands(x, w, q, k, v buffer, rows, dim int) error { + stride := 3 * dim + return firstErr( + m.r.MatMulStridedB(mb(x), mb(w), mb(q), rows, dim, dim, stride, 0), + m.r.MatMulStridedB(mb(x), mb(w), mb(k), rows, dim, dim, stride, dim), + m.r.MatMulStridedB(mb(x), mb(w), mb(v), rows, dim, dim, stride, 2*dim), + ) +} +func (m mRec) RoPE(q, inv, o buffer, seq, width, heads, hd, half, pos int, posDiv float32) error { + return m.r.RoPE(mb(q), mb(inv), mb(o), seq, width, heads, hd, half, pos, posDiv) +} +func (m mRec) RoPEAt(q, inv, o buffer, off, seq, width, heads, hd, half, pos int, posDiv float32) error { + return m.r.RoPEAt(mb(q), mb(inv), mb(o), off, seq, width, heads, hd, half, pos, posDiv) +} +func (m mRec) RoPEPair(qkv, inv buffer, seq, stride, headsQ, offQ, headsK, offK, hd, half, pos int, posDiv float32) error { + return m.r.RoPEPair(mb(qkv), mb(inv), seq, stride, headsQ, offQ, headsK, offK, hd, half, pos, posDiv) +} +func (m mRec) RoPEPairSplit(qkv, inv, q, k, v buffer, + seq, stride, headsQ, offQ, headsK, offK, hd, half, vOff, vDim, pos int, posDiv float32, +) error { + return m.r.RoPEPairSplit(mb(qkv), mb(inv), mb(q), mb(k), mb(v), + seq, stride, headsQ, offQ, headsK, offK, hd, half, vOff, vDim, pos, posDiv) +} +func (mRec) RoPEPartialPair(qkv, inv buffer, seq, stride, headsQ, offQ, headsK, offK, hd, rotaryDim, pos int, posDiv float32) error { + return fmt.Errorf("llamagpu(metal): partial-rotary RoPE not implemented (partial-rotary decoders are cuda-only for now)") +} +func (m mRec) Blit(src buffer, srcOff int, dst buffer, dstOff, n int) error { + return m.r.Blit(mb(src), srcOff, mb(dst), dstOff, n) +} +func (m mRec) Copy2D(src buffer, srcOff, srcStride int, dst buffer, dstOff, dstStride, rows, rowFloats int) error { + return m.r.Copy2D(mb(src), srcOff, srcStride, mb(dst), dstOff, dstStride, rows, rowFloats) +} +func (m mRec) Copy2DF32ToF16Pair( + kSrc buffer, kSrcOff, kSrcStride int, + vSrc buffer, vSrcOff, vSrcStride int, + kDst buffer, kDstOff, kDstStride int, + vDst buffer, vDstOff, vDstStride int, + rows, rowFloats int, +) error { + return m.r.Copy2DF32ToF16Pair( + mb(kSrc), kSrcOff, kSrcStride, + mb(vSrc), vSrcOff, vSrcStride, + mb(kDst), kDstOff, kDstStride, + mb(vDst), vDstOff, vDstStride, + rows, rowFloats, + ) +} +func (m mRec) RoPEF16KVAppend(q, k, v, inv, kCache, vCache buffer, + headsQ, headsK, hd, half, pos, cacheOff int, posDiv float32, +) error { + return m.r.RoPEF16KVAppend( + mb(q), mb(k), mb(v), mb(inv), mb(kCache), mb(vCache), + headsQ, headsK, hd, half, pos, cacheOff, posDiv, + ) +} +func (m mRec) RoPEPairF16KVAppend(qkv, inv, kCache, vCache buffer, + stride, headsQ, offQ, headsK, offK, hd, half, vOff, vDim, pos, cacheOff int, posDiv float32, +) error { + return m.r.RoPEPairF16KVAppend( + mb(qkv), mb(inv), mb(kCache), mb(vCache), + stride, headsQ, offQ, headsK, offK, hd, half, vOff, vDim, pos, cacheOff, posDiv, + ) +} +func (m mRec) MHA(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error { + return m.r.MHA(mb(q), mb(k), mb(v), mb(o), sq, sk, dm, heads, kvHeads, dk, causal, window, scale) +} +func (m mRec) MHAF16KV(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error { + return m.r.MHAF16KV(mb(q), mb(k), mb(v), mb(o), sq, sk, dm, heads, kvHeads, dk, causal, window, scale) +} +func (mRec) MHACap(q, k, v, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale, cap float32) error { + return fmt.Errorf("llamagpu(metal): attention-logit soft-cap not implemented (Gemma-2-class softcap decoders are cuda-only for now)") +} +func (mRec) MHAALiBi(q, k, v, o, slopes buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error { + return fmt.Errorf("llamagpu(metal): ALiBi attention not implemented (ALiBi decoders are cuda-only for now)") +} +func (mRec) MHABias(q, k, v, o, bias buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error { + return fmt.Errorf("llamagpu(metal): per-head bias attention not implemented (T5 encoder is cuda-only for now)") +} +func (mRec) MoEGate(logits, weights buffer, rows, e, k, raw int, scale float32) error { + return fmt.Errorf("llamagpu(metal): MoE routing not implemented (sparse-MoE decoders are cuda-only for now)") +} +func (mRec) RowAxpy(dst, src, arow buffer, rows, cols int) error { + return fmt.Errorf("llamagpu(metal): MoE combine not implemented (sparse-MoE decoders are cuda-only for now)") +} +func (mRec) MHARect(q, k, v, o buffer, sq, sk, heads, kvHeads, dqk, dv, causal, window int, scale float32) error { + return fmt.Errorf("llamagpu(metal): rectangular MHA not implemented (MLA decoders are cuda-only for now)") +} +func (mRec) SSMStep(u, delta, a, b, c, dskip, h, y buffer, d, n int) error { + return fmt.Errorf("llamagpu(metal): SSM step not implemented (Mamba decoders are cuda-only for now)") +} +func (mRec) Conv1DStep(x, w, b, state, out buffer, d, k int) error { + return fmt.Errorf("llamagpu(metal): conv1d step not implemented (Mamba decoders are cuda-only for now)") +} +func (mRec) SSDStep(x, delta, a, b, c, dskip, state, y buffer, heads, headDim, groups, n int) error { + return fmt.Errorf("llamagpu(metal): SSD step not implemented (Mamba-2 decoders are cuda-only for now)") +} +func (mRec) WKVStep(k, v, w, u, aa, bb, pp, out buffer, d int) error { + return fmt.Errorf("llamagpu(metal): WKV step not implemented (RWKV decoders are cuda-only for now)") +} +func (m mRec) Unary(x, o buffer, op int) error { return m.r.Unary(mb(x), mb(o), op) } +func (m mRec) Binary(a, b, o buffer, op int) error { + return m.r.Binary(mb(a), mb(b), mb(o), op) +} +func (m mRec) BinaryN(a, b, o buffer, op, n int) error { + return m.r.BinaryN(mb(a), mb(b), mb(o), op, n) +} +func (m mRec) QMatMulResident(x buffer, w qweight, o buffer, mm int) error { + switch w := w.(type) { + case *metal.ResidentQWeight: + return m.r.QMatMulResident(mb(x), w, mb(o), mm) + case *metal.ResidentQGroup: + return m.r.QMatMulResidentGroup(mb(x), w, mb(o), mm) + default: + return fmt.Errorf("llamagpu(metal): unsupported resident quant weight %T", w) + } +} +func (m mRec) Commit() error { return m.r.Commit() } +func (m mRec) Wait() error { return m.r.Wait() } +func (m mRec) Finish() error { return m.r.Finish() } +func (m mRec) Free() { m.r.Free() } + +func (m mProfileRec) Finish() error { + if err := m.r.Finish(); err != nil { + return err + } + p, err := m.r.Profile() + if err == nil { + *m.profile = p + } + return err +} + +func (m mProfileRec) Wait() error { + if err := m.r.Wait(); err != nil { + return err + } + p, err := m.r.Profile() + if err == nil { + *m.profile = p + } + return err +} + +func mb(b buffer) *metal.DeviceBuffer { return b.(mBuf).DeviceBuffer } + +// New uploads m's weights into metal device buffers and prepares a KV (key/value attention cache) cache up to m.Config.Ctx +// tokens (the 24× batched decode, §T404). +func New(m *nlp.Llama) (*Decoder, error) { + if !metal.Available() { + return nil, fmt.Errorf("llamagpu: no metal GPU") + } + return newDecoder(m, backendOps{ + name: string(backend.Metal), + asyncEncode: true, // metal command buffers are independent objects (§T614) + newBuffer: func(data []float32) (buffer, error) { + b, err := metal.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return mBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := metal.NewRecorder() + if err != nil { + return nil, err + } + return mRec{r}, nil + }, + uploadQWeight: metalUploadQWeight, + }) +} + +// NewGPT uploads an nlp.GPT's weights into metal device buffers for batched decoding — the +// GPT-2-style sibling of New (§T422). +func NewGPT(m *nlp.GPT) (*GPTDecoder, error) { + if !metal.Available() { + return nil, fmt.Errorf("llamagpu: no metal GPU") + } + return newGPTDecoder(m, backendOps{ + name: string(backend.Metal), + fusedF32QKV: true, + fusedBiasGELU: os.Getenv("GOAI_METAL_GPT_BIAS_GELU") != "0", + newBuffer: func(data []float32) (buffer, error) { + b, err := metal.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return mBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := metal.NewRecorder() + if err != nil { + return nil, err + } + return mRec{r}, nil + }, + }) +} + +func metalUploadQWeight(weight []byte, qt uint32, n, k int) (qweight, error) { + if qt == 3 { // GGUF Q4_1: recorder-only residency; host QuantLinear stays on faster ARM64. + return metal.UploadQWeightQ4_1(weight, n, k) + } + if qt == 16 { // GGUF IQ2_XXS: exact eight-value grids, recorder-only after the M2 host-route gate. + return metal.UploadQWeightIQ2_XXS(weight, n, k) + } + if qt == 17 { // GGUF IQ2_XS: exact 9-bit grid blocks, recorder-only after the M2 host-route gate. + return metal.UploadQWeightIQ2_XS(weight, n, k) + } + if qt == 18 { // GGUF IQ3_XXS: exact grid blocks, recorder-only after the M2 host-route gate. + return metal.UploadQWeightIQ3_XXS(weight, n, k) + } + if qt == 19 { // GGUF IQ1_S: shared exact packed ternary grid, recorder-only on M2. + return metal.UploadQWeightIQ1_S(weight, n, k) + } + if qt == 20 { // GGUF IQ4_NL: same recorder-only boundary after the M2 host-route gate. + return metal.UploadQWeightIQ4_NL(weight, n, k) + } + if qt == 21 { // GGUF IQ3_S: exact 9-bit grid blocks, recorder-only after the M2 host-route gate. + return metal.UploadQWeightIQ3_S(weight, n, k) + } + if qt == 22 { // GGUF IQ2_S: exact packed 10-bit grid blocks, recorder-only after the M2 host-route gate. + return metal.UploadQWeightIQ2_S(weight, n, k) + } + if qt == 23 { // GGUF IQ4_XS: exact 256-value nonlinear super-blocks, recorder-only on M2. + return metal.UploadQWeightIQ4_XS(weight, n, k) + } + if qt == 29 { // GGUF IQ1_M: split-f16 ternary blocks, recorder-only on M2. + return metal.UploadQWeightIQ1_M(weight, n, k) + } + if qt == 34 { // GGUF TQ1_0: base-243 ternary blocks, recorder-only on M2. + return metal.UploadQWeightTQ1_0(weight, n, k) + } + if qt == 35 { // GGUF TQ2_0: two-bit ternary blocks, recorder-only on M2. + return metal.UploadQWeightTQ2_0(weight, n, k) + } + if qt == 39 { // GGUF MXFP4: OCP E8M0/E2M1 blocks, recorder-only on M2. + return metal.UploadQWeightMXFP4(weight, n, k) + } + if qt == 41 { // GGUF Q1_0: packed binary blocks, recorder-only on M2. + return metal.UploadQWeightQ1_0(weight, n, k) + } + rw, err := metal.Backend{}.UploadQuant(weight, qt, n, k) + if err != nil { + return nil, err + } + return rw.(*metal.ResidentQWeight), nil +} + +func metalGroupQWeights(weights []qweight) (qweight, error) { + if len(weights) != 3 { + return nil, fmt.Errorf("llamagpu(metal): mixed QKV group needs 3 weights, got %d", len(weights)) + } + concrete := make([]*metal.ResidentQWeight, 3) + for i, w := range weights { + var ok bool + concrete[i], ok = w.(*metal.ResidentQWeight) + if !ok { + return nil, fmt.Errorf("llamagpu(metal): mixed QKV weight %d has type %T", i, w) + } + } + return metal.NewResidentQGroup(concrete...) +} + +// NewQuant uploads a quantized Llama (nlp.QuantizeLlama / nlp.QuantLlamaFromGGUF) with every +// projection as a device-resident quantized weight — the 4-8× smaller weights of quantization +// combined with the batched-decode speedup (§T415). Metal. +func NewQuant(m *nlp.QuantLlama) (*Decoder, error) { + return newQuantMetal(m, false) +} + +// NewQuantF16KV is NewQuant with an opt-in IEEE-binary16 K/V cache. It halves retained cache +// storage and uses half-reading Metal attention while Q/O and accumulation remain f32. The initial +// prompt still uses the established f32 FlashMM path as its K/V rows are converted into the cache. +// NewQuant and every non-Metal backend remain unchanged. The current M2-specialized path requires +// head dimension 64 and returns an explicit error for other geometries. +func NewQuantF16KV(m *nlp.QuantLlama) (*Decoder, error) { + return newQuantMetal(m, true) +} + +func newQuantMetal(m *nlp.QuantLlama, f16KV bool) (*Decoder, error) { + return newQuantMetalWithMixedQKV(m, f16KV, true) +} + +// newQuantMetalWithMixedQKV is the production constructor plus a narrow control seam for the +// trained-model promotion campaign. Public constructors always enable the grouped path; tests can +// build the otherwise-identical established separate-projection arm without environment switches. +func newQuantMetalWithMixedQKV(m *nlp.QuantLlama, f16KV, mixedQKV bool) (*Decoder, error) { + if !metal.Available() { + return nil, fmt.Errorf("llamagpu: no metal GPU") + } + ops := backendOps{ + name: string(backend.Metal), + asyncEncode: true, // metal command buffers are independent objects (§T614) + newBuffer: func(data []float32) (buffer, error) { + b, err := metal.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return mBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := metal.NewRecorder() + if err != nil { + return nil, err + } + return mRec{r}, nil + }, + uploadQWeight: metalUploadQWeight, + } + if mixedQKV { + ops.groupQWeights = metalGroupQWeights + } + if f16KV { + ops.newF16KVBuffer = func(n int) (buffer, error) { + b, err := metal.NewDeviceBufferF16Zeros(n) + if err != nil { + return nil, err + } + return mBuf{b}, nil + } + } + d, err := newQuantDecoder(m, ops) + if err != nil { + return nil, err + } + if concurrentMetalDecodeEligible(d) { + d.ops.newDecodeRecorder = func() (recorder, error) { + r, err := metal.NewConcurrentRecorder() + if err != nil { + return nil, err + } + return mRec{r}, nil + } + } + return d, nil +} + +// concurrentMetalDecodeEligible deliberately starts with the dense pre-norm Llama graph whose +// dependencies are explicit in encodeStep. Architectures with biased, parallel, post-norm, MoE, +// state-space, partial-RoPE, or fused FFN projection graphs retain the established recorder. +func concurrentMetalDecodeEligible(d *Decoder) bool { + if d == nil || !d.f16KV || d.rwkv || d.mamba || d.mamba2 || d.jamba || d.mla || + d.postNorm || d.sandwich || d.parallelRes || d.parallelTwoNorm || d.qkNorm || d.noRope || + d.lnBias || d.ffnGELU || d.ffnReLU2 || d.ffnGEGLU || d.rotaryDim != d.dk || + d.attnCap != 0 || d.aliBiSlopes != nil || d.outBias != nil { + return false + } + for i := range d.blocks { + b := &d.blocks[i] + if b.moeRouter != nil || b.wGU != nil || b.qkvBias != nil || + b.oBias != nil || b.fcBias != nil || b.projBias != nil { + return false + } + } + return true +} + +// ProfileMetalStep runs one ordinary Decoder.Step through an opt-in Metal timestamp recorder and +// returns both its logits and per-encoder GPU durations. It disables next-step pre-encoding for +// this diagnostic call so exactly one command buffer is attributed, then restores the decoder's +// production recorder factory. Any pending pre-encoded step is discarded; Decoder is already +// documented as not safe for concurrent use. +func (d *Decoder) ProfileMetalStep(token, pos, maxEvents int) ([]float32, metal.RecorderProfile, error) { + if d.ops.name != string(backend.Metal) { + return nil, metal.RecorderProfile{}, fmt.Errorf("llamagpu: ProfileMetalStep requires the Metal backend, got %s", d.ops.name) + } + if pos < 0 || pos >= d.maxLen { + return nil, metal.RecorderProfile{}, fmt.Errorf("llamagpu(%s): pos %d out of [0,%d)", d.ops.name, pos, d.maxLen) + } + if token < 0 || token >= d.v { + return nil, metal.RecorderProfile{}, fmt.Errorf("llamagpu(%s): token %d out of vocab %d", d.ops.name, token, d.v) + } + if maxEvents < 1 || maxEvents > 2048 { + return nil, metal.RecorderProfile{}, fmt.Errorf("llamagpu: ProfileMetalStep maxEvents=%d outside [1,2048]", maxEvents) + } + d.dropPending() + oldNewRecorder, oldNewDecodeRecorder, oldAsyncEncode := d.ops.newRecorder, d.ops.newDecodeRecorder, d.ops.asyncEncode + defer func() { + d.ops.newRecorder = oldNewRecorder + d.ops.newDecodeRecorder = oldNewDecodeRecorder + d.ops.asyncEncode = oldAsyncEncode + }() + var profile metal.RecorderProfile + d.ops.newRecorder = func() (recorder, error) { + r, err := metal.NewProfilingRecorder(maxEvents) + if err != nil { + return nil, err + } + return mProfileRec{mRec: mRec{r: r}, profile: &profile}, nil + } + d.ops.newDecodeRecorder = d.ops.newRecorder + d.ops.asyncEncode = false + logits, err := d.Step(token, pos) + if err != nil { + return nil, metal.RecorderProfile{}, err + } + return logits, profile, nil +} diff --git a/checks/testdata/ps6136-owner/model-api.go.txt b/checks/testdata/ps6136-owner/model-api.go.txt new file mode 100644 index 00000000..97509983 --- /dev/null +++ b/checks/testdata/ps6136-owner/model-api.go.txt @@ -0,0 +1,74 @@ +package nlp +import("github.com/jxsl13/goai/tensor";"github.com/jxsl13/goai/nn";"github.com/jxsl13/goai/backend") +type Block struct{LN1 *nn.LayerNorm; LN2 *nn.LayerNorm; Attn *MHA; W1 *tensor.Tensor; B1 *tensor.Tensor; W2 *tensor.Tensor; B2 *tensor.Tensor} +type Cohere struct{Config CohereConfig; TokEmb *tensor.Tensor; Blocks []*CohereBlock; Norm *nn.LayerNorm; Out *tensor.Tensor} +type CohereBlock struct{InputNorm *nn.LayerNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; FFN *nn.SwiGLU} +type CohereConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; HeadDim int; Layers int; Hidden int; Eps float64; RopeBase float64; LogitScale float64} +type DeepSeekV2 struct{Config DeepSeekV2Config; TokEmb *tensor.Tensor; Blocks []*DeepSeekV2Block; FinalNorm *nn.RMSNorm; LmHead *tensor.Tensor} +type DeepSeekV2Block struct{InputNorm *nn.RMSNorm; WqA *tensor.Tensor; QANorm *nn.RMSNorm; WqB *tensor.Tensor; WkvA *tensor.Tensor; KvANorm *nn.RMSNorm; WkvB *tensor.Tensor; Wo *tensor.Tensor; PostAttnNorm *nn.RMSNorm; Dense *nn.SwiGLU; MoE *nn.DeepSeekMoE} +type DeepSeekV2Config struct{Vocab int; Ctx int; Dim int; Heads int; QLoraRank int; KVLoraRank int; QKNope int; QKRope int; VHead int; Layers int; FFN int; Eps float64; RopeBase float64; SoftmaxScale float64; FirstKDense int; NRoutedExperts int; NSharedExperts int; TopK int; MoEHidden int; RoutedScale float64} +type Falcon struct{Config FalconConfig; TokEmb *tensor.Tensor; Blocks []*FalconBlock; FinalNorm *nn.LayerNorm; Out *tensor.Tensor} +type FalconBlock struct{InputNorm *nn.LayerNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; Wh *tensor.Tensor; Wout *tensor.Tensor} +type FalconConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; Layers int; Hidden int; Eps float64; RopeBase float64} +type GPT struct{Config GPTConfig; TokEmb *tensor.Tensor; PosEmb *tensor.Tensor; Blocks []*Block; LNf *nn.LayerNorm; Head *tensor.Tensor; LayerBackends []backend.Name} +type GPTConfig struct{Vocab int; Ctx int; Dim int; Heads int; Layers int; Eps float64} +type GPTNeoX struct{Config GPTNeoXConfig; TokEmb *tensor.Tensor; Blocks []*GPTNeoXBlock; FinalNorm *nn.LayerNorm; Out *tensor.Tensor} +type GPTNeoXBlock struct{InputNorm *nn.LayerNorm; PostAttnNorm *nn.LayerNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; Bq *tensor.Tensor; Bk *tensor.Tensor; Bv *tensor.Tensor; Bo *tensor.Tensor; Wh *tensor.Tensor; Bh *tensor.Tensor; Wout *tensor.Tensor; Bout *tensor.Tensor} +type GPTNeoXConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; Layers int; Hidden int; Eps float64; RopeBase float64; RotaryPct float64; RotaryDim int} +type Gemma struct{Config GemmaConfig; TokEmb *tensor.Tensor; Blocks []*GemmaBlock; FinalNorm *nn.RMSNorm; outT *tensor.Tensor} +type Gemma2 struct{Config Gemma2Config; TokEmb *tensor.Tensor; Blocks []*Gemma2Block; FinalNorm *nn.RMSNorm; outT *tensor.Tensor} +type Gemma2Block struct{InputNorm *nn.RMSNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; PostAttnNorm *nn.RMSNorm; PreFFNNorm *nn.RMSNorm; FFN *nn.GLU; PostFFNNorm *nn.RMSNorm} +type Gemma2Config struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; HeadDim int; Layers int; FFN int; Eps float64; RopeBase float64; QueryPreAttnScalar float64; AttnLogitCap float64; FinalLogitCap float64} +type GemmaBlock struct{AttnNorm *nn.RMSNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; FFNNorm *nn.RMSNorm; FFN *nn.GLU} +type GemmaConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; HeadDim int; Layers int; FFN int; Eps float64; RopeBase float64} +type GraniteMoE struct{Config GraniteMoeConfig; TokEmb *tensor.Tensor; Blocks []*GraniteMoeBlock; Norm *nn.RMSNorm; Out *tensor.Tensor} +type GraniteMoeBlock struct{AttnNorm *nn.RMSNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; FFNNorm *nn.RMSNorm; MoE *nn.SparseMoE} +type GraniteMoeConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; Layers int; Hidden int; Experts int; TopK int; Eps float64; RopeBase float64; EmbeddingMult float64; AttentionMult float64; ResidualMult float64; LogitsScale float64} +type Jamba struct{Config JambaConfig; TokEmb *tensor.Tensor; Layers []*JambaLayer; Norm *nn.RMSNorm; Out *tensor.Tensor} +type JambaAttention struct{Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor} +type JambaConfig struct{Vocab int; Dim int; Heads int; KVHeads int; Layers int; TopK int; Eps float64} +type JambaLayer struct{InputNorm *nn.RMSNorm; PreFFNorm *nn.RMSNorm; Attn *JambaAttention; Mamba *JambaMixer; MoE *JambaMoE; Dense *nn.SwiGLU} +type JambaMixer struct{Block *nn.MambaBlock; DtNorm *nn.RMSNorm; BNorm *nn.RMSNorm; CNorm *nn.RMSNorm} +type JambaMoE struct{Router *tensor.Tensor; Experts []*nn.SwiGLU; TopK int} +type Llama struct{Config LlamaConfig; TokEmb *tensor.Tensor; Blocks []*LlamaBlock; Norm *nn.RMSNorm; Out *tensor.Tensor; LayerBackends []backend.Name} +type LlamaBlock struct{AttnNorm *nn.RMSNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; Bq *tensor.Tensor; Bk *tensor.Tensor; Bv *tensor.Tensor; QNorm *nn.RMSNorm; KNorm *nn.RMSNorm; FFNNorm *nn.RMSNorm; FFN *nn.SwiGLU} +type LlamaConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; Layers int; Hidden int; Eps float64; RopeBase float64; EmbeddingMult float64; AttentionMult float64; ResidualMult float64; LogitsScale float64} +type MHA struct{Heads int; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; Causal bool; LoRA map[string]*nn.LoRALinear; Bias map[string]*tensor.Tensor; Mask *tensor.Tensor} +type MPT struct{Config MPTConfig; TokEmb *tensor.Tensor; Blocks []*MPTBlock; FinalNorm *nn.LayerNorm; Out *tensor.Tensor} +type MPTBlock struct{Norm1 *nn.LayerNorm; Norm2 *nn.LayerNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; Wup *tensor.Tensor; Wdown *tensor.Tensor} +type MPTConfig struct{Vocab int; Ctx int; Dim int; Heads int; Layers int; Hidden int; Eps float64} +type Mamba struct{Config MambaConfig; Embed *tensor.Tensor; Layers []MambaLayer; Norm *nn.RMSNorm; Head *tensor.Tensor} +type Mamba2 struct{Config Mamba2Config; Embed *tensor.Tensor; Layers []Mamba2Layer; Norm *nn.RMSNorm; Head *tensor.Tensor} +type Mamba2Config struct{DModel int; NumHeads int; HeadDim int; NGroups int; N int; DConv int; Intermediate int; Layers int; Vocab int; Eps float64} +type Mamba2Layer struct{Norm *nn.RMSNorm; Mixer *Mamba2Mixer} +type Mamba2Mixer struct{InProj *tensor.Tensor; ConvW *tensor.Tensor; ConvB *tensor.Tensor; ALog *tensor.Tensor; D *tensor.Tensor; DtBias *tensor.Tensor; NormW *tensor.Tensor; OutProj *tensor.Tensor; inProjT *tensor.Tensor; outProjT *tensor.Tensor; DModel int; NumHeads int; HeadDim int; NGroups int; N int; DConv int; Intermediate int; ConvDim int; Eps float64} +type MambaConfig struct{DModel int; N int; DConv int; Expand int; DtRank int; Layers int; Vocab int; Eps float64} +type MambaLayer struct{Norm *nn.RMSNorm; Mixer *nn.MambaBlock} +type Mixtral struct{Config MixtralConfig; TokEmb *tensor.Tensor; Blocks []*MixtralBlock; Norm *nn.RMSNorm; Out *tensor.Tensor} +type MixtralBlock struct{AttnNorm *nn.RMSNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; QNorm *nn.RMSNorm; KNorm *nn.RMSNorm; FFNNorm *nn.RMSNorm; MoE *nn.SparseMoE} +type MixtralConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; Layers int; Hidden int; Experts int; TopK int; Eps float64; RopeBase float64} +type Nemotron struct{Config NemotronConfig; TokEmb *tensor.Tensor; Blocks []*NemotronBlock; Norm *nn.LayerNorm; Out *tensor.Tensor} +type NemotronBlock struct{InputNorm *nn.LayerNorm; PostAttnNorm *nn.LayerNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; Wup *tensor.Tensor; Wdown *tensor.Tensor} +type NemotronConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; HeadDim int; Layers int; Hidden int; Eps float64; RopeBase float64; RotaryDim int; RotaryPct float64} +type OLMo2 struct{Config OLMo2Config; TokEmb *tensor.Tensor; Blocks []*OLMo2Block; Norm *nn.RMSNorm; Out *tensor.Tensor} +type OLMo2Block struct{Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; QNorm *nn.RMSNorm; KNorm *nn.RMSNorm; PostAttnNorm *nn.RMSNorm; FFN *nn.SwiGLU; PostFFNNorm *nn.RMSNorm} +type OLMo2Config struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; HeadDim int; Layers int; Hidden int; Eps float64; RopeBase float64} +type OLMoE struct{Config OLMoEConfig; TokEmb *tensor.Tensor; Blocks []*OLMoEBlock; Norm *nn.RMSNorm; Out *tensor.Tensor} +type OLMoEBlock struct{AttnNorm *nn.RMSNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; QNorm *nn.RMSNorm; KNorm *nn.RMSNorm; FFNNorm *nn.RMSNorm; MoE *nn.SparseMoE} +type OLMoEConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; HeadDim int; Layers int; Hidden int; Experts int; TopK int; Eps float64; RopeBase float64} +type Phi struct{Config PhiConfig; TokEmb *tensor.Tensor; Blocks []*PhiBlock; FinalNorm *nn.LayerNorm; Out *tensor.Tensor; OutBias *tensor.Tensor} +type PhiBlock struct{InputNorm *nn.LayerNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wdense *tensor.Tensor; Bq *tensor.Tensor; Bk *tensor.Tensor; Bv *tensor.Tensor; Bdense *tensor.Tensor; Wfc1 *tensor.Tensor; Bfc1 *tensor.Tensor; Wfc2 *tensor.Tensor; Bfc2 *tensor.Tensor} +type PhiConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; HeadDim int; Layers int; Hidden int; Eps float64; RopeBase float64; RotaryDim int; RotaryPct float64} +type QuantBlock struct{AttnNorm *nn.RMSNorm; Wq *nn.QuantLinear; Wk *nn.QuantLinear; Wv *nn.QuantLinear; Wo *nn.QuantLinear; Bq *tensor.Tensor; Bk *tensor.Tensor; Bv *tensor.Tensor; QNorm *nn.RMSNorm; KNorm *nn.RMSNorm; FFNNorm *nn.RMSNorm; FFN *nn.QuantSwiGLU} +type QuantLlama struct{Config LlamaConfig; TokEmb *tensor.Tensor; Blocks []*QuantBlock; Norm *nn.RMSNorm; Out *nn.QuantLinear} +type Qwen2MoE struct{Config Qwen2MoeConfig; TokEmb *tensor.Tensor; Blocks []*Qwen2MoeBlock; Norm *nn.RMSNorm; Out *tensor.Tensor} +type Qwen2MoeBlock struct{AttnNorm *nn.RMSNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; Bq *tensor.Tensor; Bk *tensor.Tensor; Bv *tensor.Tensor; FFNNorm *nn.RMSNorm; MoE *nn.SparseMoE; Shared *nn.SwiGLU; SharedGate *tensor.Tensor} +type Qwen2MoeConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; Layers int; Hidden int; SharedHidden int; Experts int; TopK int; Eps float64; RopeBase float64} +type RWKV struct{Config RWKVConfig; Embed *tensor.Tensor; PreLN *nn.LayerNorm; Blocks []*nn.RWKVBlock; LNOut *nn.LayerNorm; Head *tensor.Tensor} +type RWKVConfig struct{Vocab int; Dim int; Hidden int; Layers int; Eps float64} +type StableLM struct{Config StableLMConfig; TokEmb *tensor.Tensor; Blocks []*StableLMBlock; Norm *nn.LayerNorm; Out *tensor.Tensor} +type StableLMBlock struct{InputNorm *nn.LayerNorm; PostAttnNorm *nn.LayerNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; FFN *nn.SwiGLU} +type StableLMConfig struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; HeadDim int; Layers int; Hidden int; Eps float64; RopeBase float64; RotaryDim int; RotaryPct float64} +type StarCoder2 struct{Config StarCoder2Config; TokEmb *tensor.Tensor; Blocks []*StarCoder2Block; Norm *nn.LayerNorm; Out *tensor.Tensor} +type StarCoder2Block struct{InputNorm *nn.LayerNorm; PostAttnNorm *nn.LayerNorm; Wq *tensor.Tensor; Wk *tensor.Tensor; Wv *tensor.Tensor; Wo *tensor.Tensor; Bq *tensor.Tensor; Bk *tensor.Tensor; Bv *tensor.Tensor; Bo *tensor.Tensor; Wfc *tensor.Tensor; Bfc *tensor.Tensor; Wproj *tensor.Tensor; Bproj *tensor.Tensor} +type StarCoder2Config struct{Vocab int; Ctx int; Dim int; Heads int; KVHeads int; HeadDim int; Layers int; Hidden int; Eps float64; RopeBase float64} diff --git a/checks/testdata/ps6136-owner/nn-api.go.txt b/checks/testdata/ps6136-owner/nn-api.go.txt new file mode 100644 index 00000000..faa14ea0 --- /dev/null +++ b/checks/testdata/ps6136-owner/nn-api.go.txt @@ -0,0 +1,15 @@ +package nn +import("sync";"github.com/jxsl13/goai/tensor";"github.com/jxsl13/goai/backend";"github.com/jxsl13/goai/format/gguf") +type DeepSeekMoE struct{Shared []*SwiGLU; Routed *SparseMoE} +type GLU struct{Wgate *tensor.Tensor; Wup *tensor.Tensor; Wdown *tensor.Tensor; act backend.Op; name string} +type LayerNorm struct{Gamma *tensor.Tensor; Beta *tensor.Tensor; Eps float64} +type Linear struct{W *tensor.Tensor; B *tensor.Tensor} +type LoRALinear struct{W *tensor.Tensor; A *tensor.Tensor; B *tensor.Tensor; R int; Alpha float64} +type LossFreeBalancer struct{Bias []float64; Rate float64} +type MambaBlock struct{InX *Linear; InZ *Linear; ConvW *tensor.Tensor; ConvB *tensor.Tensor; DtLow *Linear; BProj *Linear; CProj *Linear; DtProj *Linear; ALog *tensor.Tensor; Dskip *tensor.Tensor; OutProj *Linear; DModel int; DInner int; DConv int; N int; DtRank int} +type QuantLinear struct{Weight []byte; QT gguf.QuantType; In int; Out int; resOnce sync.Once; res backend.ResidentWeight} +type QuantSwiGLU struct{Gate *QuantLinear; Up *QuantLinear; Down *QuantLinear} +type RMSNorm struct{Gamma *tensor.Tensor; Eps float64} +type RWKVBlock struct{Dim int; LN1 *LayerNorm; LN2 *LayerNorm; MuR *tensor.Tensor; MuK *tensor.Tensor; MuV *tensor.Tensor; Wr *Linear; Wk *Linear; Wv *Linear; Wo *Linear; WLog *tensor.Tensor; U *tensor.Tensor; CMuR *tensor.Tensor; CMuK *tensor.Tensor; CWr *Linear; CWk *Linear; CWv *Linear} +type SparseMoE struct{Router *Linear; Experts []*SwiGLU; TopK int; balancer *LossFreeBalancer; loadCount []int} +type SwiGLU struct{Wgate *tensor.Tensor; Wup *tensor.Tensor; Wdown *tensor.Tensor} diff --git a/checks/testdata/ps6136-owner/sampling_fastpath.go.txt b/checks/testdata/ps6136-owner/sampling_fastpath.go.txt new file mode 100644 index 00000000..a22360ca --- /dev/null +++ b/checks/testdata/ps6136-owner/sampling_fastpath.go.txt @@ -0,0 +1,135 @@ +package llamagpu + +import ( + "os" + "sort" + + "github.com/jxsl13/goai/nlp" + "github.com/jxsl13/goai/tensor" +) + +// deviceTopKer is implemented by a device logits buffer that can return its k highest (index, value) +// pairs selected at the device-resident boundary. CUDA computes them on the GPU; Metal scans its +// coherent Shared MTLBuffer directly on Apple Silicon and materializes only O(k) result data. Vulkan +// does not implement it, so the sampling fast-path type-asserts this and falls back to the full-vocab +// host path when it is absent. Keeping this an interface keeps decoder.go backend-neutral. +type deviceTopKer interface { + // TopKN returns the k highest (index, value) pairs over the FIRST n elements (n = vocab; the decode + // logits occupy [0,n) of a buffer that may be over-allocated for prefill/batch). + TopKN(n, k int) ([]int32, []float32, error) +} + +// deviceTopPer extends deviceTopKer with the full-vocab softmax stats and a whole-buffer download needed +// to resolve a pure-top-p nucleus on-device (see fastTopPSampler / nlp.SampleTopPFromCandidates): the +// nucleus is Z-dependent, so the candidates alone don't determine it — SoftmaxStatsN supplies (max, Zexp) +// and ToHost is the overflow fallback. The CUDA logits buffer (cBuf, which embeds *cuda.DeviceF32) +// satisfies it for free; other backends fall back to the full-vocab host path. +type deviceTopPer interface { + TopKN(n, k int) ([]int32, []float32, error) + SoftmaxStatsN(n int, temperature float64) (maxLogit, zexp float64, err error) + ToHost() (*tensor.Tensor, error) +} + +// deviceTopKSamplingDisabled honors the backend-neutral switch and the original CUDA-named switch +// for compatibility with existing benchmark and fallback-parity campaigns. +func deviceTopKSamplingDisabled() bool { + return os.Getenv("GOAI_DEVICE_TOPK_SAMPLE") == "0" || os.Getenv("GOAI_CUDA_TOPK_SAMPLE") == "0" +} + +// fastTopKSampler reports (sampler, K) when s is a penalty-free *nlp.Sampler with TopK>0 whose selection +// is confined to the top-TopK — so a device TopK(K>TopK) superset lets the CPU draw EXACTLY on the K +// candidates (see sampleTopKCandidates), avoiding the whole-vocab logits D2H + CPU sort per token. +// Returns (nil, 0) for any other sampler (penalties on, TopK off/too large, a non-*Sampler impl, or +// GOAI_DEVICE_TOPK_SAMPLE=0) → the flexible full-vocab host fallback. The original +// GOAI_CUDA_TOPK_SAMPLE=0 spelling remains a compatibility alias. +func fastTopKSampler(s nlp.TokenSampler, vocab int) (*nlp.Sampler, int) { + if deviceTopKSamplingDisabled() { + return nil, 0 + } + sp, ok := s.(*nlp.Sampler) + if !ok { + return nil, 0 + } + // Penalties rewrite arbitrary tokens' logits from history, so a raw-logit top-k is not the effective + // top-k → require them off (RepeatPenalty 0 or 1 are both no-ops). XTC/MinP/Typical/TopP are fine: + // they only ever shrink the post-TopK set, which stays ⊆ the K candidates. + if !(sp.RepeatPenalty == 0 || sp.RepeatPenalty == 1) || sp.FreqPenalty != 0 || sp.PresencePenalty != 0 || sp.DRYMultiplier != 0 { + return nil, 0 + } + if sp.TopK <= 0 || sp.TopK > 240 { // K = TopK+16 must stay within the cu_topk cap (256) + return nil, 0 + } + k := sp.TopK + 16 // margin so a tie at rank TopK can't push a real top-TopK token out of the K set + if k > vocab { + k = vocab + } + return sp, k +} + +// toppCandidateK is the device-TopK width used to resolve a pure-top-p nucleus on-device. 64 covers a +// realistic confident-decode nucleus (typically 10–40 tokens at top-p 0.9–0.95) while keeping the O(n·K) +// cu_topk_f32 cheap — K=256 measured ~6× slower (3.1ms vs 0.5ms at 128k), erasing the win. A nucleus that +// exceeds 64 overflows to the host fallback (cheaply predicted by the C/Zexp guard before the TopK runs). +const toppCandidateK = 64 + +// fastTopPSampler reports (sampler, C) when s is a penalty-free PURE top-p *nlp.Sampler — TopP∈(0,1) +// with every OTHER truncation filter off (no TopK/MinP/Typical/Eta/Epsilon/TopNSigma, no penalties, +// no XTC) and Temperature>0 — so a device TopK(C) superset lets SampleTopPFromCandidates reconstruct +// the exact full-vocab nucleus and draw, avoiding the whole-vocab logits D2H + host sort per token. +// The nucleus is Z-dependent (unlike top-k), so the caller must also supply the full-vocab softmax +// stats (maxLogit, Zexp); see SampleTopPFromCandidates. Returns (nil,0) for any richer sampler → +// the flexible full-vocab host fallback. Disabled by GOAI_DEVICE_TOPK_SAMPLE=0. +func fastTopPSampler(s nlp.TokenSampler, vocab int) (*nlp.Sampler, int) { + if deviceTopKSamplingDisabled() { + return nil, 0 + } + sp, ok := s.(*nlp.Sampler) + if !ok { + return nil, 0 + } + // Pure top-p only: any other set-shrinking filter changes the kept set (and Typical/MinP have their + // own Z-dependence), and penalties rewrite arbitrary logits — none are handled by the top-p nucleus + // reconstruction, so require them all off. (Typical==1 is "off" per its field docs.) + if sp.Temperature <= 0 || sp.TopP <= 0 || sp.TopP >= 1 { + return nil, 0 + } + if sp.TopK > 0 || sp.MinP != 0 || sp.Epsilon != 0 || sp.Eta != 0 || sp.TopNSigma > 0 { + return nil, 0 + } + if sp.Typical != 0 && sp.Typical != 1 { + return nil, 0 + } + if !(sp.RepeatPenalty == 0 || sp.RepeatPenalty == 1) || sp.FreqPenalty != 0 || sp.PresencePenalty != 0 || sp.DRYMultiplier != 0 { + return nil, 0 + } + if sp.XTCProbability != 0 { + return nil, 0 + } + c := toppCandidateK + if c > vocab { + c = vocab + } + return sp, c +} + +// sampleTopKCandidates draws exactly as the full-vocab sampler would, but over just the K device top-k +// candidates: it feeds them to sp.Sample in ASCENDING VOCAB-INDEX order so the multinomial's cumulative +// scan visits the selected tokens in the same order (and consumes the same rng) as the full-vocab path, +// then maps the chosen local index back to its vocab id. Exact for the penalty-free TopK case (K⊇top-k). +func sampleTopKCandidates(sp *nlp.Sampler, gi []int32, gv []float32) int { + k := len(gi) + order := make([]int, k) + for j := range order { + order[j] = j + } + sort.Slice(order, func(a, b int) bool { return gi[order[a]] < gi[order[b]] }) + vals := make([]float64, k) + for j, o := range order { + vals[j] = float64(gv[o]) + } + local := sp.Sample(vals) + if local < 0 || local >= k { + local = 0 // degenerate (no positive mass) — the sampler already fell back to argmax on vals + } + return int(gi[order[local]]) +} diff --git a/checks/testdata/ps6136-owner/vulkan.go.txt b/checks/testdata/ps6136-owner/vulkan.go.txt new file mode 100644 index 00000000..7664d9cb --- /dev/null +++ b/checks/testdata/ps6136-owner/vulkan.go.txt @@ -0,0 +1,184 @@ +//go:build vulkan && cgo + +package llamagpu + +import ( + "fmt" + + "github.com/jxsl13/goai/backend" + "github.com/jxsl13/goai/backend/vulkan" + "github.com/jxsl13/goai/nlp" +) + +// vulkan adapter (§T409): thin assertions from the backend-agnostic buffer/recorder interfaces to +// the concrete vulkan.DeviceBuffer/vulkan.Recorder — the same API as metal by construction (§T408). + +type vBuf struct{ *vulkan.DeviceBuffer } + +type vRec struct{ r *vulkan.Recorder } + +func (v vRec) RMSNorm(x, g, o buffer, rows, dim int, eps float32) error { + return v.r.RMSNorm(vb(x), vb(g), vb(o), rows, dim, eps) +} +func (v vRec) LayerNorm(x, g, b, o buffer, rows, dim int, eps float32) error { + return v.r.LayerNorm(vb(x), vb(g), vb(b), vb(o), rows, dim, eps) +} +func (v vRec) AddBias(x, b, o buffer, rows, n int) error { + return v.r.AddBias(vb(x), vb(b), vb(o), rows, n) +} +func (v vRec) MatMul(a, b, c buffer, m, k, n int) error { + return v.r.MatMul(vb(a), vb(b), vb(c), m, k, n) +} +func (v vRec) MatMulAcc(a, b, c buffer, m, k, n int) error { + return v.r.MatMulAcc(vb(a), vb(b), vb(c), m, k, n) +} +func (v vRec) RoPE(q, inv, o buffer, seq, width, heads, hd, half, pos int, posDiv float32) error { + return v.r.RoPE(vb(q), vb(inv), vb(o), seq, width, heads, hd, half, pos, posDiv) +} +func (v vRec) RoPEAt(q, inv, o buffer, off, seq, width, heads, hd, half, pos int, posDiv float32) error { + return v.r.RoPEAt(vb(q), vb(inv), vb(o), off, seq, width, heads, hd, half, pos, posDiv) +} +func (v vRec) RoPEPair(qkv, inv buffer, seq, stride, headsQ, offQ, headsK, offK, hd, half, pos int, posDiv float32) error { + return v.r.RoPEPair(vb(qkv), vb(inv), seq, stride, headsQ, offQ, headsK, offK, hd, half, pos, posDiv) +} +func (vRec) RoPEPartialPair(qkv, inv buffer, seq, stride, headsQ, offQ, headsK, offK, hd, rotaryDim, pos int, posDiv float32) error { + return fmt.Errorf("llamagpu(vulkan): partial-rotary RoPE not implemented (partial-rotary decoders are cuda-only for now)") +} +func (v vRec) Blit(src buffer, srcOff int, dst buffer, dstOff, n int) error { + return v.r.Blit(vb(src), srcOff, vb(dst), dstOff, n) +} +func (v vRec) Copy2D(src buffer, srcOff, srcStride int, dst buffer, dstOff, dstStride, rows, rowFloats int) error { + return v.r.Copy2D(vb(src), srcOff, srcStride, vb(dst), dstOff, dstStride, rows, rowFloats) +} +func (v vRec) MHA(q, k, va, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error { + return v.r.MHA(vb(q), vb(k), vb(va), vb(o), sq, sk, dm, heads, kvHeads, dk, causal, window, scale) +} +func (vRec) MHACap(q, k, va, o buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale, cap float32) error { + return fmt.Errorf("llamagpu(vulkan): attention-logit soft-cap not implemented (Gemma-2-class softcap decoders are cuda-only for now)") +} +func (vRec) MHAALiBi(q, k, va, o, slopes buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error { + return fmt.Errorf("llamagpu(vulkan): ALiBi attention not implemented (ALiBi decoders are cuda-only for now)") +} +func (vRec) MHABias(q, k, va, o, bias buffer, sq, sk, dm, heads, kvHeads, dk, causal, window int, scale float32) error { + return fmt.Errorf("llamagpu(vulkan): per-head bias attention not implemented (T5 encoder is cuda-only for now)") +} +func (vRec) MoEGate(logits, weights buffer, rows, e, k, raw int, scale float32) error { + return fmt.Errorf("llamagpu(vulkan): MoE routing not implemented (sparse-MoE decoders are cuda-only for now)") +} +func (vRec) RowAxpy(dst, src, arow buffer, rows, cols int) error { + return fmt.Errorf("llamagpu(vulkan): MoE combine not implemented (sparse-MoE decoders are cuda-only for now)") +} +func (vRec) MHARect(q, k, va, o buffer, sq, sk, heads, kvHeads, dqk, dv, causal, window int, scale float32) error { + return fmt.Errorf("llamagpu(vulkan): rectangular MHA not implemented (MLA decoders are cuda-only for now)") +} +func (vRec) SSMStep(u, delta, a, b, c, dskip, h, y buffer, d, n int) error { + return fmt.Errorf("llamagpu(vulkan): SSM step not implemented (Mamba decoders are cuda-only for now)") +} +func (vRec) Conv1DStep(x, w, b, state, out buffer, d, k int) error { + return fmt.Errorf("llamagpu(vulkan): conv1d step not implemented (Mamba decoders are cuda-only for now)") +} +func (vRec) SSDStep(x, delta, a, b, c, dskip, state, y buffer, heads, headDim, groups, n int) error { + return fmt.Errorf("llamagpu(vulkan): SSD step not implemented (Mamba-2 decoders are cuda-only for now)") +} +func (vRec) WKVStep(k, v, w, u, aa, bb, pp, out buffer, d int) error { + return fmt.Errorf("llamagpu(vulkan): WKV step not implemented (RWKV decoders are cuda-only for now)") +} +func (v vRec) Unary(x, o buffer, op int) error { return v.r.Unary(vb(x), vb(o), op) } +func (v vRec) Binary(a, b, o buffer, op int) error { + return v.r.Binary(vb(a), vb(b), vb(o), op) +} +func (v vRec) QMatMulResident(x buffer, w qweight, o buffer, m int) error { + return v.r.QMatMulResident(vb(x), w.(*vulkan.ResidentQWeight), vb(o), m) +} +func (v vRec) Commit() error { return v.r.Commit() } +func (v vRec) Wait() error { return v.r.Wait() } +func (v vRec) Finish() error { return v.r.Finish() } +func (v vRec) Free() { v.r.Free() } + +func vb(b buffer) *vulkan.DeviceBuffer { return b.(vBuf).DeviceBuffer } + +// NewVulkan uploads m's weights into vulkan device buffers and prepares a KV cache up to +// m.Config.Ctx tokens — the vulkan variant of the batched decode (decode is dispatch-bound on +// vulkan too; §T390 measured 6.15× batched-vs-per-op at the recorder level). +func NewVulkan(m *nlp.Llama) (*Decoder, error) { + if !vulkan.Available() { + return nil, fmt.Errorf("llamagpu: no vulkan GPU") + } + return newDecoder(m, backendOps{ + name: string(backend.Vulkan), + newBuffer: func(data []float32) (buffer, error) { + b, err := vulkan.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return vBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := vulkan.NewRecorder() + if err != nil { + return nil, err + } + return vRec{r}, nil + }, + uploadQWeight: vulkanUploadQWeight, + }) +} + +// NewGPTVulkan uploads an nlp.GPT's weights into vulkan device buffers for batched decoding — +// the vulkan variant of NewGPT (§T422). +func NewGPTVulkan(m *nlp.GPT) (*GPTDecoder, error) { + if !vulkan.Available() { + return nil, fmt.Errorf("llamagpu: no vulkan GPU") + } + return newGPTDecoder(m, backendOps{ + name: string(backend.Vulkan), + newBuffer: func(data []float32) (buffer, error) { + b, err := vulkan.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return vBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := vulkan.NewRecorder() + if err != nil { + return nil, err + } + return vRec{r}, nil + }, + }) +} + +func vulkanUploadQWeight(weight []byte, qt uint32, n, k int) (qweight, error) { + rw, err := vulkan.Backend{}.UploadQuant(weight, qt, n, k) + if err != nil { + return nil, err + } + return rw.(*vulkan.ResidentQWeight), nil +} + +// NewQuantVulkan uploads a quantized Llama with every projection as a device-resident quantized +// weight — the vulkan variant of NewQuant (§T415). +func NewQuantVulkan(m *nlp.QuantLlama) (*Decoder, error) { + if !vulkan.Available() { + return nil, fmt.Errorf("llamagpu: no vulkan GPU") + } + return newQuantDecoder(m, backendOps{ + name: string(backend.Vulkan), + newBuffer: func(data []float32) (buffer, error) { + b, err := vulkan.NewDeviceBufferF32(data) + if err != nil { + return nil, err + } + return vBuf{b}, nil + }, + newRecorder: func() (recorder, error) { + r, err := vulkan.NewRecorder() + if err != nil { + return nil, err + } + return vRec{r}, nil + }, + uploadQWeight: vulkanUploadQWeight, + }) +} diff --git a/checks/testdata/src/ps6136/gpt.go b/checks/testdata/src/ps6136/gpt.go new file mode 100644 index 00000000..06f5463c --- /dev/null +++ b/checks/testdata/src/ps6136/gpt.go @@ -0,0 +1,389 @@ +package llamagpu + +import ( + "fmt" + "math" + + "github.com/jxsl13/goai/nlp" + "github.com/jxsl13/goai/tensor" +) + +// gptBlock is one pre-LN GPT transformer block on the device: LayerNorm gains/biases, the four +// attention projections, the biased GELU MLP, and the block's KV cache. +type gptBlock struct { + g1, b1, g2, b2 buffer // LN1/LN2 gamma+beta [d] + wq, wk, wv, wo linear // attention projections [d,d], no bias + wqkv f32Linear // fused attention projection [d,3d] + w1, w2 linear // FFN up [d,4d] / down [4d,d] + fb1, fb2 buffer // FFN biases [4d] / [d] + kC, vC buffer // KV caches [maxLen*d] + ffn int // FFN width, fixed when the block is uploaded +} + +// GPTDecoder is the batched-decode engine for nlp.GPT models (§T422) — the GPT-2-style sibling of +// [Decoder]: LayerNorm instead of RMSNorm, learned positional embeddings instead of RoPE, a biased +// GELU MLP instead of SwiGLU. Each Step records the whole layer stack into one command buffer. +// Not safe for concurrent use; Release when done. +type GPTDecoder struct { + ops backendOps + d, h, dk, v, maxLen int + eps, scale float32 + + blocks []gptBlock + lnfG, lnfB *bufSlot + head linear + dx, xn, xn2 *bufSlot + q, k, v_, qkv *bufSlot + attn, hid *bufSlot + logits *bufSlot + tokEmb, posEmb *tensor.Tensor // host gather sources + all []buffer +} + +type f32QKVBandsRecorder interface { + F32QKVBands(x, w, q, k, v buffer, rows, dim int) error +} + +type biasGELURecorder interface { + BiasGELU(x, b, o buffer, rows, n int) error +} + +func (d *GPTDecoder) recordAttentionResidual(r recorder, b gptBlock, rows int) error { + return b.wo.recordAdd(r, d.attn.b, nil, d.dx.b, rows, d.d) +} + +func (d *GPTDecoder) recordFFNResidual(r recorder, b gptBlock, rows int) error { + return firstErr( + b.w2.recordAdd(r, d.hid.b, nil, d.dx.b, rows, d.d), + r.AddBias(d.dx.b, b.fb2, d.dx.b, rows, d.d), + ) +} + +func (d *GPTDecoder) recordBiasGELU(r recorder, b gptBlock, rows int) error { + if d.ops.fusedBiasGELU { + if fused, ok := r.(biasGELURecorder); ok { + return fused.BiasGELU(d.hid.b, b.fb1, d.hid.b, rows, b.ffn) + } + } + return firstErr( + r.AddBias(d.hid.b, b.fb1, d.hid.b, rows, b.ffn), + r.Unary(d.hid.b, d.hid.b, unaryGELU), + ) +} + +// newGPTDecoder uploads m's weights via ops and prepares per-block KV caches up to Ctx. +func newGPTDecoder(m *nlp.GPT, ops backendOps) (*GPTDecoder, error) { + cfg := m.Config + d := &GPTDecoder{ + ops: ops, + d: cfg.Dim, h: cfg.Heads, v: cfg.Vocab, maxLen: cfg.Ctx, + eps: float32(cfg.Eps), tokEmb: m.TokEmb, posEmb: m.PosEmb, + } + if d.h <= 0 || d.d%d.h != 0 { + return nil, fmt.Errorf("llamagpu(%s): gpt dim %d not divisible by heads %d", ops.name, d.d, d.h) + } + d.dk = d.d / d.h + d.scale = float32(1.0 / math.Sqrt(float64(d.dk))) + + var err error + mk := func(data []float32) *bufSlot { + if err != nil { + return &bufSlot{} + } + b, e := ops.newBuffer(data) + if e != nil { + err = e + return &bufSlot{} + } + d.all = append(d.all, b) + return &bufSlot{b: b} + } + lin := func(w *tensor.Tensor) linear { + in, out := w.Shape()[0], w.Shape()[1] + return f32Linear{w: mk(flat2D(w)).b, k: in, n: out} + } + fusedQKV := func(wq, wk, wv *tensor.Tensor) f32Linear { + in := wq.Shape()[0] + nq, nk, nv := wq.Shape()[1], wk.Shape()[1], wv.Shape()[1] + fq, fk, fv := flat2D(wq), flat2D(wk), flat2D(wv) + stride := nq + nk + nv + weight := make([]float32, in*stride) + for row := range in { + dst := weight[row*stride : (row+1)*stride] + copy(dst[:nq], fq[row*nq:(row+1)*nq]) + copy(dst[nq:nq+nk], fk[row*nk:(row+1)*nk]) + copy(dst[nq+nk:], fv[row*nv:(row+1)*nv]) + } + return f32Linear{w: mk(weight).b, k: in, n: stride} + } + for _, b := range m.Blocks { + gb := gptBlock{ + g1: mk(flat1D(b.LN1.Gamma)).b, b1: mk(flat1D(b.LN1.Beta)).b, + g2: mk(flat1D(b.LN2.Gamma)).b, b2: mk(flat1D(b.LN2.Beta)).b, + wo: lin(b.Attn.Wo), w1: lin(b.W1), w2: lin(b.W2), + fb1: mk(flat1D(b.B1)).b, fb2: mk(flat1D(b.B2)).b, + kC: mk(make([]float32, d.maxLen*d.d)).b, vC: mk(make([]float32, d.maxLen*d.d)).b, + ffn: b.W1.Shape()[1], + } + if ops.fusedF32QKV { + gb.wqkv = fusedQKV(b.Attn.Wq, b.Attn.Wk, b.Attn.Wv) + } else { + gb.wq, gb.wk, gb.wv = lin(b.Attn.Wq), lin(b.Attn.Wk), lin(b.Attn.Wv) + } + d.blocks = append(d.blocks, gb) + } + d.lnfG = mk(flat1D(m.LNf.Gamma)) + d.lnfB = mk(flat1D(m.LNf.Beta)) + d.head = lin(m.Head) + ffn := 4 * d.d + if len(m.Blocks) > 0 { + ffn = m.Blocks[0].W1.Shape()[1] + } + c := d.maxLen // scratch sized for multi-token StepN (§T423) + d.dx = mk(make([]float32, c*d.d)) + d.xn = mk(make([]float32, c*d.d)) + d.xn2 = mk(make([]float32, c*d.d)) + d.q = mk(make([]float32, c*d.d)) + d.k = mk(make([]float32, c*d.d)) + d.v_ = mk(make([]float32, c*d.d)) + if ops.fusedF32QKV { + // Large prefills address bands of the single fused weight directly. Only rows below the + // strided-band crossover need a grouped-output scratch, bounding it to 63 rows. + d.qkv = mk(make([]float32, min(c, 63)*3*d.d)) + } + d.attn = mk(make([]float32, c*d.d)) + d.hid = mk(make([]float32, c*ffn)) + d.logits = mk(make([]float32, c*d.v)) // want "reviewed common one-row output policy retains" + if err != nil { + d.Release() + return nil, err + } + return d, nil +} + +// Step advances the decoder by one token at absolute position pos (== cache length before the +// call), recording the whole layer stack + head into one command buffer, and returns the [vocab] +// logits. The token embedding and learned positional embedding are summed host-side. +func (d *GPTDecoder) Step(token, pos int) ([]float32, error) { + if pos < 0 || pos >= d.maxLen { + return nil, fmt.Errorf("llamagpu(%s): gpt pos %d out of [0,%d)", d.ops.name, pos, d.maxLen) + } + if token < 0 || token >= d.v { + return nil, fmt.Errorf("llamagpu(%s): gpt token %d out of vocab %d", d.ops.name, token, d.v) + } + x := embedRow(d.tokEmb, token, d.d) + pe := embedRow(d.posEmb, pos, d.d) + for i := range x { + x[i] += pe[i] + } + if err := d.dx.b.UploadF32(x); err != nil { + return nil, err + } + r, err := d.ops.newRecorder() + if err != nil { + return nil, err + } + D, H, dk := d.d, d.h, d.dk + for _, b := range d.blocks { + e := r.LayerNorm(d.dx.b, b.g1, b.b1, d.xn.b, 1, D, d.eps) + qBuf := d.q.b + if d.ops.fusedF32QKV { + qBuf = d.qkv.b + e = firstErr(e, + b.wqkv.record(r, d.xn.b, d.qkv.b, 1), + r.Blit(d.qkv.b, D, b.kC, pos*D, D), + r.Blit(d.qkv.b, 2*D, b.vC, pos*D, D), + ) + } else { + e = firstErr(e, + b.wq.record(r, d.xn.b, d.q.b, 1), + b.wk.record(r, d.xn.b, d.k.b, 1), + b.wv.record(r, d.xn.b, d.v_.b, 1), + r.Blit(d.k.b, 0, b.kC, pos*D, D), + r.Blit(d.v_.b, 0, b.vC, pos*D, D), + ) + } + e = firstErr(e, + r.MHA(qBuf, b.kC, b.vC, d.attn.b, 1, pos+1, D, H, H, dk, 1, 0, d.scale), + d.recordAttentionResidual(r, b, 1), + r.LayerNorm(d.dx.b, b.g2, b.b2, d.xn2.b, 1, D, d.eps), + b.w1.record(r, d.xn2.b, d.hid.b, 1), + d.recordBiasGELU(r, b, 1), + d.recordFFNResidual(r, b, 1), + ) + if e != nil { + r.Free() + return nil, e + } + } + if e := firstErr( + r.LayerNorm(d.dx.b, d.lnfG.b, d.lnfB.b, d.xn.b, 1, D, d.eps), + d.head.record(r, d.xn.b, d.logits.b, 1), + r.Finish(), + ); e != nil { + r.Free() + return nil, e + } + r.Free() + out := make([]float32, d.v) + if err := d.logits.b.DownloadF32(out); err != nil { + return nil, err + } + return out, nil +} + +// StepN advances the decoder by k tokens at absolute positions pos..pos+k-1 in ONE recorded step +// (§T423, the GPT sibling of Decoder.StepN): the layer stack runs over [k,·] rows with causal +// attention against the growing cache, biases broadcast via the AddBias record op, and the k KV +// rows are appended. Returns the [k,vocab] logits. +func (d *GPTDecoder) StepN(tokens []int, pos int) ([]float32, error) { + return d.gptStepN(tokens, pos, false) +} + +// StepNLast is StepN projecting only the FINAL prompt row's logits (LM head over 1 row, download of +// vocab not k·vocab); the transformer body still runs all k rows so the KV cache and the returned +// [vocab] row are bit-identical to StepN's last row. For generation/KV-seeding prefills that read at +// most the last row — mirrors Decoder.StepNLast. +func (d *GPTDecoder) StepNLast(tokens []int, pos int) ([]float32, error) { + return d.gptStepN(tokens, pos, true) +} + +func (d *GPTDecoder) gptStepN(tokens []int, pos int, lastOnly bool) ([]float32, error) { + k := len(tokens) + if k == 0 { + return nil, fmt.Errorf("llamagpu(%s): gpt StepN needs ≥1 token", d.ops.name) + } + if pos < 0 || pos+k > d.maxLen { + return nil, fmt.Errorf("llamagpu(%s): gpt StepN [%d,%d) out of [0,%d)", d.ops.name, pos, pos+k, d.maxLen) + } + host := make([]float32, k*d.d) + for i, tok := range tokens { + if tok < 0 || tok >= d.v { + return nil, fmt.Errorf("llamagpu(%s): gpt token %d out of vocab %d", d.ops.name, tok, d.v) + } + row := host[i*d.d : (i+1)*d.d] + copy(row, embedRow(d.tokEmb, tok, d.d)) + pe := embedRow(d.posEmb, pos+i, d.d) + for j := range row { + row[j] += pe[j] + } + } + if err := d.dx.b.UploadF32(host); err != nil { + return nil, err + } + r, err := d.ops.newRecorder() + if err != nil { + return nil, err + } + D, H, dk := d.d, d.h, d.dk + for _, b := range d.blocks { + e := r.LayerNorm(d.dx.b, b.g1, b.b1, d.xn.b, k, D, d.eps) + if d.ops.fusedF32QKV { + if bands, ok := r.(f32QKVBandsRecorder); ok && k >= 64 { + e = firstErr(e, bands.F32QKVBands(d.xn.b, b.wqkv.w, d.q.b, d.k.b, d.v_.b, k, D)) + } else { + e = firstErr(e, + b.wqkv.record(r, d.xn.b, d.qkv.b, k), + r.Copy2D(d.qkv.b, 0, 3*D, d.q.b, 0, D, k, D), + r.Copy2D(d.qkv.b, D, 3*D, d.k.b, 0, D, k, D), + r.Copy2D(d.qkv.b, 2*D, 3*D, d.v_.b, 0, D, k, D), + ) + } + } else { + e = firstErr(e, + b.wq.record(r, d.xn.b, d.q.b, k), + b.wk.record(r, d.xn.b, d.k.b, k), + b.wv.record(r, d.xn.b, d.v_.b, k), + ) + } + e = firstErr(e, + r.Blit(d.k.b, 0, b.kC, pos*D, k*D), + r.Blit(d.v_.b, 0, b.vC, pos*D, k*D), + r.MHA(d.q.b, b.kC, b.vC, d.attn.b, k, pos+k, D, H, H, dk, 1, 0, d.scale), + d.recordAttentionResidual(r, b, k), + r.LayerNorm(d.dx.b, b.g2, b.b2, d.xn2.b, k, D, d.eps), + b.w1.record(r, d.xn2.b, d.hid.b, k), + d.recordBiasGELU(r, b, k), + d.recordFFNResidual(r, b, k), + ) + if e != nil { + r.Free() + return nil, e + } + } + rows := k + if lastOnly { + // Only the final row's logits are wanted: move its residual to row 0 so the final LN + LM head + // project ONE row and the download shrinks from k·vocab to vocab (bit-identical — row 0 then + // holds exactly what row k-1 would have produced; a no-op when k==1). + if e := r.Blit(d.dx.b, (k-1)*D, d.dx.b, 0, D); e != nil { + r.Free() + return nil, e + } + rows = 1 + } + if e := firstErr( + r.LayerNorm(d.dx.b, d.lnfG.b, d.lnfB.b, d.xn.b, rows, D, d.eps), + d.head.record(r, d.xn.b, d.logits.b, rows), + r.Finish(), + ); e != nil { + r.Free() + return nil, e + } + r.Free() + out := make([]float32, rows*d.v) + if err := d.logits.b.DownloadF32(out); err != nil { + return nil, err + } + return out, nil +} + +// Generate prefills the whole prompt in ONE recorded StepN (§T423), then samples up to maxNew +// tokens (bounded by Ctx), returning prompt+generated ids. +func (d *GPTDecoder) Generate(prompt []int, maxNew int, s nlp.TokenSampler) ([]int, error) { + if len(prompt) == 0 { + return nil, fmt.Errorf("llamagpu(%s): gpt Generate needs a non-empty prompt", d.ops.name) + } + out := append([]int(nil), prompt...) + // generation samples only the post-prompt row, so project just the last prompt row's logits + // (StepNLast) instead of the full [len(prompt), vocab] LM head + download. + logits, err := d.StepNLast(prompt, 0) + if err != nil { + return nil, err + } + pos := len(prompt) + buf := make([]float64, d.v) + for range maxNew { + if pos >= d.maxLen { + break + } + for i, x := range logits { + buf[i] = float64(x) + } + next := s.SampleWithHistory(buf, out) + out = append(out, next) + l, err := d.Step(next, pos) + if err != nil { + return nil, err + } + logits = l + pos++ + } + return out, nil +} + +// Vocab returns the model's vocabulary size. +func (d *GPTDecoder) Vocab() int { return d.v } + +// Ctx returns the model's maximum context length (the KV-cache capacity). +func (d *GPTDecoder) Ctx() int { return d.maxLen } + +// Release frees all device buffers. +func (d *GPTDecoder) Release() { + for _, b := range d.all { + if b != nil { + b.Release() + } + } + d.all = nil +} diff --git a/config/config.go b/config/config.go index 79b525b2..9f7c099a 100644 --- a/config/config.go +++ b/config/config.go @@ -252,6 +252,7 @@ type Config struct { // ActiveBoundFallbackContracts bind optional explicit-bound interface APIs // and their reviewed capacity-wide fallback semantics for PS6135. ActiveBoundFallbackContracts []ActiveBoundFallbackContract `json:"activeBoundFallbackContracts,omitempty" yaml:"activeBoundFallbackContracts"` + OutputWorkspaceContracts []OutputWorkspaceContract `json:"outputWorkspaceContracts,omitempty" yaml:"outputWorkspaceContracts"` // NativeSnapshotReuseContracts separate source acquisition/ownership roles // from reviewed opaque native lifetime and repeated-extraction policy. @@ -2361,6 +2362,7 @@ type Sets struct { InPlaceFusionContracts []InPlaceFusionContract BoundedScratchFlowContracts []BoundedScratchFlowContract ActiveBoundFallbackContracts []ActiveBoundFallbackContract + OutputWorkspaceContracts []OutputWorkspaceContract NativeSnapshotReuseContracts []NativeSnapshotReuseContract ReceiverStagingContracts []ReceiverStagingContract ReusableOneShotWrapperContracts []ReusableOneShotWrapperContract @@ -2437,6 +2439,7 @@ func (c Config) Compile() Sets { //perfscan:ignore PS3106 one startup call; keep InPlaceFusionContracts: slices.Clone(c.InPlaceFusionContracts), BoundedScratchFlowContracts: cloneBoundedScratchFlowContracts(c.BoundedScratchFlowContracts), ActiveBoundFallbackContracts: cloneActiveBoundFallbackContracts(c.ActiveBoundFallbackContracts), + OutputWorkspaceContracts: cloneOutputWorkspaceContracts(c.OutputWorkspaceContracts), NativeSnapshotReuseContracts: cloneNativeSnapshotReuseContracts(c.NativeSnapshotReuseContracts), ReceiverStagingContracts: slices.Clone(c.ReceiverStagingContracts), ReusableOneShotWrapperContracts: cloneReusableOneShotWrapperContracts(c.ReusableOneShotWrapperContracts), diff --git a/config/output_workspace.go b/config/output_workspace.go new file mode 100644 index 00000000..77216cb4 --- /dev/null +++ b/config/output_workspace.go @@ -0,0 +1,256 @@ +package config + +import "slices" + +// OutputWorkspaceLeaf names an exact typed API and its reviewed buffer role. +// A negative element argument denotes a source-derived host descriptor length, +// or the explicitly inventoried whole-capacity transfer, never an unknown bound. +type OutputWorkspaceLeaf struct { + Method string `json:"method" yaml:"method"` + Kind string `json:"kind" yaml:"kind"` + BufferArguments []int `json:"bufferArguments" yaml:"bufferArguments"` + RowsArgument int `json:"rowsArgument" yaml:"rowsArgument"` + WidthArgument int `json:"widthArgument" yaml:"widthArgument"` + HostArgument int `json:"hostArgument" yaml:"hostArgument"` + SemanticsReviewed bool `json:"semanticsReviewed" yaml:"semanticsReviewed"` + Implementations []string `json:"implementations" yaml:"implementations"` + HostStorageMethod string `json:"hostStorageMethod,omitempty" yaml:"hostStorageMethod"` + HostFloat32Method string `json:"hostFloat32Method,omitempty" yaml:"hostFloat32Method"` + SelectedCapabilityAbsenceAllowed bool `json:"selectedCapabilityAbsenceAllowed,omitempty" yaml:"selectedCapabilityAbsenceAllowed"` +} + +// OutputWorkspaceContract separates workload and opaque backend semantics +// from mandatory source proof of constructor, exact buffer and argument flow. +type OutputWorkspaceContract struct { + BackendOpsField string `json:"backendOpsField" yaml:"backendOpsField"` + BackendRecorderField string `json:"backendRecorderField" yaml:"backendRecorderField"` + NativeRecorderFactory string `json:"nativeRecorderFactory" yaml:"nativeRecorderFactory"` + NativeRecorderType string `json:"nativeRecorderType" yaml:"nativeRecorderType"` + RecorderWrapperType string `json:"recorderWrapperType" yaml:"recorderWrapperType"` + RecorderNativeField string `json:"recorderNativeField" yaml:"recorderNativeField"` + BufferWrapperType string `json:"bufferWrapperType" yaml:"bufferWrapperType"` + BufferNativeField string `json:"bufferNativeField" yaml:"bufferNativeField"` + BufferBridge string `json:"bufferBridge" yaml:"bufferBridge"` + ProjectorType string `json:"projectorType" yaml:"projectorType"` + ProjectorWeightField string `json:"projectorWeightField" yaml:"projectorWeightField"` + ProjectorInnerField string `json:"projectorInnerField" yaml:"projectorInnerField"` + ProjectorWidthField string `json:"projectorWidthField" yaml:"projectorWidthField"` + ProjectorRecordMethod string `json:"projectorRecordMethod" yaml:"projectorRecordMethod"` + RecorderProjectionMethod string `json:"recorderProjectionMethod" yaml:"recorderProjectionMethod"` + NativeProjectionMethod string `json:"nativeProjectionMethod" yaml:"nativeProjectionMethod"` + NativeAccessSemanticsReviewed bool `json:"nativeAccessSemanticsReviewed" yaml:"nativeAccessSemanticsReviewed"` + SecondaryRecorderField string `json:"secondaryRecorderField,omitempty" yaml:"secondaryRecorderField"` + AsyncRecorderField string `json:"asyncRecorderField,omitempty" yaml:"asyncRecorderField"` + RecorderOverrideMethod string `json:"recorderOverrideMethod,omitempty" yaml:"recorderOverrideMethod"` + NativeOverrideFactory string `json:"nativeOverrideFactory,omitempty" yaml:"nativeOverrideFactory"` + RecorderOverrideWrapperType string `json:"recorderOverrideWrapperType,omitempty" yaml:"recorderOverrideWrapperType"` + RecorderOverrideEmbedField string `json:"recorderOverrideEmbedField,omitempty" yaml:"recorderOverrideEmbedField"` + RecorderOverrideMetadataField string `json:"recorderOverrideMetadataField,omitempty" yaml:"recorderOverrideMetadataField"` + DropPendingMethod string `json:"dropPendingMethod,omitempty" yaml:"dropPendingMethod"` + OwnerType string `json:"ownerType" yaml:"ownerType"` + RetainedBufferReleaseMethod string `json:"retainedBufferReleaseMethod" yaml:"retainedBufferReleaseMethod"` + NativeBufferReleaseMethod string `json:"nativeBufferReleaseMethod" yaml:"nativeBufferReleaseMethod"` + NativeLifetimeSemanticsReviewed bool `json:"nativeLifetimeSemanticsReviewed" yaml:"nativeLifetimeSemanticsReviewed"` + ModelType string `json:"modelType" yaml:"modelType"` + ModelConfigField string `json:"modelConfigField" yaml:"modelConfigField"` + ModelHeadField string `json:"modelHeadField" yaml:"modelHeadField"` + ConfigMaximumRowsField string `json:"configMaximumRowsField" yaml:"configMaximumRowsField"` + ConfigWidthField string `json:"configWidthField" yaml:"configWidthField"` + HeadShapeMethod string `json:"headShapeMethod" yaml:"headShapeMethod"` + WorkspaceField string `json:"workspaceField" yaml:"workspaceField"` + SlotBufferField string `json:"slotBufferField" yaml:"slotBufferField"` + MaximumRowsField string `json:"maximumRowsField" yaml:"maximumRowsField"` + WidthField string `json:"widthField" yaml:"widthField"` + ProjectorField string `json:"projectorField" yaml:"projectorField"` + RetainedListField string `json:"retainedListField" yaml:"retainedListField"` + AllocationFunction string `json:"allocationFunction" yaml:"allocationFunction"` + ConstructorEntries []string `json:"constructorEntries" yaml:"constructorEntries"` + FactoryFunction string `json:"factoryFunction,omitempty" yaml:"factoryFunction"` + BackendAllocatorField string `json:"backendAllocatorField" yaml:"backendAllocatorField"` + BackendAllocators []string `json:"backendAllocators" yaml:"backendAllocators"` + ReleaseMethod string `json:"releaseMethod" yaml:"releaseMethod"` + CommonMethods []string `json:"commonMethods" yaml:"commonMethods"` + BulkMethod string `json:"bulkMethod" yaml:"bulkMethod"` + Leaves []OutputWorkspaceLeaf `json:"leaves" yaml:"leaves"` + ProfileMaximumRows int64 `json:"profileMaximumRows,omitempty" yaml:"profileMaximumRows"` + ProfileWidth int64 `json:"profileWidth,omitempty" yaml:"profileWidth"` + DominantOneRowPolicyReviewed bool `json:"dominantOneRowPolicyReviewed" yaml:"dominantOneRowPolicyReviewed"` + ExplicitRareBulkPolicyReviewed bool `json:"explicitRareBulkPolicyReviewed" yaml:"explicitRareBulkPolicyReviewed"` + HardRealtimeNoGrowth bool `json:"hardRealtimeNoGrowth,omitempty" yaml:"hardRealtimeNoGrowth"` + FreshFloat32OwnershipReviewed bool `json:"freshFloat32OwnershipReviewed" yaml:"freshFloat32OwnershipReviewed"` + ModelProjectorWidthReviewed bool `json:"modelProjectorWidthReviewed" yaml:"modelProjectorWidthReviewed"` + SequentialLifecycleReviewed bool `json:"sequentialLifecycleReviewed" yaml:"sequentialLifecycleReviewed"` + CheckedShapeArithmeticReviewed bool `json:"checkedShapeArithmeticReviewed" yaml:"checkedShapeArithmeticReviewed"` + AllProviderBuildPathsReviewed bool `json:"allProviderBuildPathsReviewed" yaml:"allProviderBuildPathsReviewed"` + ExternalObservationsReviewed bool `json:"externalObservationsReviewed" yaml:"externalObservationsReviewed"` +} + +func (c *OutputWorkspaceContract) Valid() bool { + if !psTopKMethodIDValid(c.RetainedBufferReleaseMethod) || !psTopKMethodIDValid(c.NativeBufferReleaseMethod) || !c.NativeLifetimeSemanticsReviewed { + return false + } + if c.SecondaryRecorderField != "" && !psTopKIdentifierValid(c.SecondaryRecorderField) { + return false + } + if c.RecorderOverrideMethod != "" { + if !psTopKMethodIDValid(c.RecorderOverrideMethod) || !psTopKMethodIDValid(c.DropPendingMethod) || !psTopKFunctionIDValid(c.NativeOverrideFactory) || !psTopKFunctionIDValid(c.RecorderOverrideWrapperType) || !psTopKIdentifierValid(c.RecorderOverrideEmbedField) || !psTopKIdentifierValid(c.RecorderOverrideMetadataField) || !psTopKIdentifierValid(c.AsyncRecorderField) || c.SecondaryRecorderField == "" { + return false + } + } else if c.NativeOverrideFactory != "" || c.RecorderOverrideWrapperType != "" || c.RecorderOverrideEmbedField != "" || c.RecorderOverrideMetadataField != "" || c.AsyncRecorderField != "" || c.DropPendingMethod != "" { + return false + } + for _, identity := range []string{c.NativeRecorderType, c.RecorderWrapperType, c.BufferWrapperType, c.ProjectorType} { + if !psTopKFunctionIDValid(identity) { + return false + } + } + for _, identity := range []string{c.NativeRecorderFactory, c.BufferBridge} { + if !psTopKFunctionIDValid(identity) { + return false + } + } + for _, identity := range []string{c.ProjectorRecordMethod, c.RecorderProjectionMethod, c.NativeProjectionMethod} { + if !psTopKMethodIDValid(identity) { + return false + } + } + for _, field := range []string{c.BackendOpsField, c.BackendRecorderField, c.RecorderNativeField, c.BufferNativeField, c.ProjectorWeightField, c.ProjectorInnerField, c.ProjectorWidthField} { + if !psTopKIdentifierValid(field) { + return false + } + } + if !c.NativeAccessSemanticsReviewed || c.BackendRecorderField == c.BackendAllocatorField || c.ProjectorWeightField == c.ProjectorInnerField || c.ProjectorWeightField == c.ProjectorWidthField || c.ProjectorInnerField == c.ProjectorWidthField { + return false + } + if !psTopKFunctionIDValid(c.OwnerType) || !psTopKFunctionIDValid(c.ModelType) || !psTopKMethodIDValid(c.HeadShapeMethod) || !ps6109CallableIDValid(c.AllocationFunction) || + !ps6109CallableIDValid(c.FactoryFunction) || + !psTopKMethodIDValid(c.ReleaseMethod) || !psTopKMethodIDValid(c.BulkMethod) || + len(c.CommonMethods) == 0 || len(c.Leaves) == 0 || len(c.BackendAllocators) == 0 || len(c.ConstructorEntries) == 0 || + !c.DominantOneRowPolicyReviewed || !c.ExplicitRareBulkPolicyReviewed || c.HardRealtimeNoGrowth || + !c.FreshFloat32OwnershipReviewed || !c.ModelProjectorWidthReviewed || !c.SequentialLifecycleReviewed || + !c.CheckedShapeArithmeticReviewed || !c.AllProviderBuildPathsReviewed || !c.ExternalObservationsReviewed || + c.ProfileMaximumRows < 0 || c.ProfileWidth < 0 || (c.ProfileMaximumRows == 0) != (c.ProfileWidth == 0) { + return false + } + fields := []string{c.WorkspaceField, c.MaximumRowsField, c.WidthField, c.ProjectorField, c.RetainedListField} + seen := make(map[string]bool, len(fields)) + for _, field := range fields { + if !psTopKIdentifierValid(field) || seen[field] { + return false + } + seen[field] = true + } + if !psTopKIdentifierValid(c.SlotBufferField) || !psTopKIdentifierValid(c.BackendAllocatorField) { + return false + } + for _, field := range []string{c.ModelConfigField, c.ModelHeadField, c.ConfigMaximumRowsField, c.ConfigWidthField} { + if !psTopKIdentifierValid(field) { + return false + } + } + if c.ModelConfigField == c.ModelHeadField || c.ConfigMaximumRowsField == c.ConfigWidthField { + return false + } + methods := map[string]bool{c.BulkMethod: true} + entries := make(map[string]bool, len(c.ConstructorEntries)) + for _, entry := range c.ConstructorEntries { + if !ps6109CallableIDValid(entry) || entries[entry] { + return false + } + entries[entry] = true + } + allocators := make(map[string]bool, len(c.BackendAllocators)) + for _, allocator := range c.BackendAllocators { + if !ps6109CallableIDValid(allocator) || allocators[allocator] { + return false + } + allocators[allocator] = true + } + for _, method := range c.CommonMethods { + if !psTopKMethodIDValid(method) || methods[method] { + return false + } + methods[method] = true + } + leaves := make(map[string]bool, len(c.Leaves)) + for index := range c.Leaves { + leaf := &c.Leaves[index] + if !psTopKMethodIDValid(leaf.Method) || leaves[leaf.Method] || !leaf.SemanticsReviewed || + len(leaf.BufferArguments) == 0 || leaf.RowsArgument < -1 || leaf.WidthArgument < -1 || leaf.HostArgument < -1 || (len(leaf.Implementations) == 0 && !leaf.SelectedCapabilityAbsenceAllowed) || (leaf.SelectedCapabilityAbsenceAllowed && leaf.Kind != "prefix" && leaf.Kind != "capacity-transfer") { + return false + } + leaves[leaf.Method] = true + // Unresolved user formal indices/list lengths are arbitrary and sparse. + // A map keeps duplicate validation linear without inventing an ABI bound. + roles := make(map[int]bool, len(leaf.BufferArguments)) //perfscan:ignore PS3083 + for _, role := range leaf.BufferArguments { + if role < -1 || roles[role] { //perfscan:ignore PS3003 + return false + } + roles[role] = true //perfscan:ignore PS3003 + } + implementations := make(map[string]bool, len(leaf.Implementations)) + for _, implementation := range leaf.Implementations { + if !psTopKMethodIDValid(implementation) || implementations[implementation] { + return false + } + implementations[implementation] = true + } + switch leaf.Kind { + case "projection", "row-width": + if roles[-1] || leaf.RowsArgument < 0 || leaf.HostArgument != -1 || //perfscan:ignore PS3003 + (leaf.Kind == "row-width" && leaf.WidthArgument < 0) || + (leaf.Kind == "projection" && leaf.WidthArgument != -1) { + return false + } + case "download": + if len(roles) != 1 || !roles[-1] || leaf.RowsArgument != -1 || leaf.WidthArgument != -1 || leaf.HostArgument < 0 { + return false + } + case "prefix": + if len(roles) != 1 || !roles[-1] || leaf.RowsArgument != -1 || leaf.WidthArgument < 0 || leaf.HostArgument != -1 { + return false + } + case "capacity-transfer": + if len(roles) != 1 || !roles[-1] || leaf.RowsArgument != -1 || leaf.WidthArgument != -1 || leaf.HostArgument != -1 || + !psTopKMethodIDValid(leaf.HostStorageMethod) || !psTopKMethodIDValid(leaf.HostFloat32Method) { + return false + } + default: + return false + } + if leaf.Kind != "capacity-transfer" && (leaf.HostStorageMethod != "" || leaf.HostFloat32Method != "") { + return false + } + } + return true +} + +func UsableOutputWorkspaceContractCount(contracts []OutputWorkspaceContract) int { + counts := make(map[string]int, len(contracts)) + for i := range contracts { + counts[contracts[i].OwnerType+"."+contracts[i].WorkspaceField]++ + } + n := 0 + for i := range contracts { + if contracts[i].Valid() && counts[contracts[i].OwnerType+"."+contracts[i].WorkspaceField] == 1 { + n++ + } + } + return n +} + +func cloneOutputWorkspaceContracts(contracts []OutputWorkspaceContract) []OutputWorkspaceContract { + result := slices.Clone(contracts) + for i := range result { + result[i].CommonMethods = slices.Clone(result[i].CommonMethods) + result[i].Leaves = slices.Clone(result[i].Leaves) + result[i].BackendAllocators = slices.Clone(result[i].BackendAllocators) + result[i].ConstructorEntries = slices.Clone(result[i].ConstructorEntries) + for j := range result[i].Leaves { + result[i].Leaves[j].Implementations = slices.Clone(result[i].Leaves[j].Implementations) + result[i].Leaves[j].BufferArguments = slices.Clone(result[i].Leaves[j].BufferArguments) + } + } + return result +} diff --git a/config/output_workspace_test.go b/config/output_workspace_test.go new file mode 100644 index 00000000..9c1d3121 --- /dev/null +++ b/config/output_workspace_test.go @@ -0,0 +1,81 @@ +package config + +import "testing" + +func TestOutputWorkspaceLargeSparseRoles(t *testing.T) { + t.Parallel() + c := outputWorkspaceConfigFixture() + const count = 32768 + c.Leaves[0].BufferArguments = make([]int, count) + for index := range c.Leaves[0].BufferArguments { + c.Leaves[0].BufferArguments[index] = index*1024 + 7 + } + if !c.Valid() { + t.Fatal("unique unresolved sparse formals rejected before typed ABI resolution") + } + c.Leaves[0].BufferArguments[count-1] = c.Leaves[0].BufferArguments[0] + if c.Valid() { + t.Fatal("large sparse role duplicate accepted") + } +} + +func outputWorkspaceConfigFixture() OutputWorkspaceContract { + return OutputWorkspaceContract{ + OwnerType: "project.Owner", ModelType: "model.Model", ModelConfigField: "Config", ModelHeadField: "Head", ConfigMaximumRowsField: "Rows", ConfigWidthField: "Width", HeadShapeMethod: "tensor.Tensor.Shape", + WorkspaceField: "output", SlotBufferField: "b", MaximumRowsField: "max", WidthField: "width", ProjectorField: "head", RetainedListField: "all", AllocationFunction: "project.allocate", FactoryFunction: "project.construct", ConstructorEntries: []string{"project.New"}, + BackendOpsField: "ops", BackendAllocatorField: "allocate", BackendAllocators: []string{"native.NewBuffer"}, BackendRecorderField: "record", NativeRecorderFactory: "native.NewRecorder", NativeRecorderType: "native.Recorder", RecorderWrapperType: "project.Rec", RecorderNativeField: "r", BufferWrapperType: "project.Buf", BufferNativeField: "Buffer", BufferBridge: "project.bridge", + ProjectorType: "project.Head", ProjectorWeightField: "w", ProjectorInnerField: "k", ProjectorWidthField: "n", ProjectorRecordMethod: "project.Head.record", RecorderProjectionMethod: "project.Recorder.Project", NativeProjectionMethod: "native.Recorder.Project", + ReleaseMethod: "project.Owner.Release", RetainedBufferReleaseMethod: "project.Buffer.Release", NativeBufferReleaseMethod: "native.Buffer.Release", CommonMethods: []string{"project.Owner.Step"}, BulkMethod: "project.Owner.Bulk", + Leaves: []OutputWorkspaceLeaf{{Method: "project.Head.record", Kind: "projection", BufferArguments: []int{2}, RowsArgument: 3, WidthArgument: -1, HostArgument: -1, SemanticsReviewed: true, Implementations: []string{"project.Head.record"}}}, + DominantOneRowPolicyReviewed: true, ExplicitRareBulkPolicyReviewed: true, FreshFloat32OwnershipReviewed: true, ModelProjectorWidthReviewed: true, SequentialLifecycleReviewed: true, CheckedShapeArithmeticReviewed: true, AllProviderBuildPathsReviewed: true, ExternalObservationsReviewed: true, NativeAccessSemanticsReviewed: true, NativeLifetimeSemanticsReviewed: true, + } +} + +func TestOutputWorkspaceContractsValidityAndClone(t *testing.T) { + t.Parallel() + original := outputWorkspaceConfigFixture() + if !original.Valid() || UsableOutputWorkspaceContractCount([]OutputWorkspaceContract{original}) != 1 { + t.Fatal("complete contract invalid") + } + for _, test := range []struct { + name string + change func(*OutputWorkspaceContract) + }{ + {"hard_rt", func(c *OutputWorkspaceContract) { c.HardRealtimeNoGrowth = true }}, + {"missing_lifetime", func(c *OutputWorkspaceContract) { c.NativeLifetimeSemanticsReviewed = false }}, + {"missing_factory", func(c *OutputWorkspaceContract) { c.FactoryFunction = "" }}, + {"negative_profile", func(c *OutputWorkspaceContract) { c.ProfileMaximumRows = -1 }}, + {"partial_profile", func(c *OutputWorkspaceContract) { c.ProfileWidth = 1 }}, + {"duplicate_entry", func(c *OutputWorkspaceContract) { + c.ConstructorEntries = append(c.ConstructorEntries, c.ConstructorEntries[0]) + }}, + {"overlapping_dimensions", func(c *OutputWorkspaceContract) { c.ProjectorInnerField = c.ProjectorWidthField }}, + {"wrong_absence_kind", func(c *OutputWorkspaceContract) { c.Leaves[0].SelectedCapabilityAbsenceAllowed = true }}, + {"unknown_kind", func(c *OutputWorkspaceContract) { c.Leaves[0].Kind = "unknown" }}, + {"duplicate_buffer_role", func(c *OutputWorkspaceContract) { c.Leaves[0].BufferArguments = []int{2, 2} }}, + {"partial_profile_override", func(c *OutputWorkspaceContract) { c.RecorderOverrideMethod = "project.Owner.Profile" }}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + c := cloneOutputWorkspaceContracts([]OutputWorkspaceContract{original})[0] + test.change(&c) + if c.Valid() || UsableOutputWorkspaceContractCount([]OutputWorkspaceContract{c}) != 0 { + t.Fatal("contradictory contract usable") + } + }) + } + duplicate := []OutputWorkspaceContract{original, original} + if UsableOutputWorkspaceContractCount(duplicate) != 0 { + t.Fatal("duplicate site usable") + } + cfg := Config{OutputWorkspaceContracts: []OutputWorkspaceContract{original}} + sets := cfg.Compile() + sets.OutputWorkspaceContracts[0].ConstructorEntries[0] = "changed.New" + sets.OutputWorkspaceContracts[0].CommonMethods[0] = "changed.Owner.Step" + sets.OutputWorkspaceContracts[0].BackendAllocators[0] = "changed.Alloc" + sets.OutputWorkspaceContracts[0].Leaves[0].BufferArguments[0] = 0 + sets.OutputWorkspaceContracts[0].Leaves[0].Implementations[0] = "changed.Head.record" + if cfg.OutputWorkspaceContracts[0].ConstructorEntries[0] != "project.New" || cfg.OutputWorkspaceContracts[0].CommonMethods[0] != "project.Owner.Step" || cfg.OutputWorkspaceContracts[0].BackendAllocators[0] != "native.NewBuffer" || cfg.OutputWorkspaceContracts[0].Leaves[0].BufferArguments[0] != 2 || cfg.OutputWorkspaceContracts[0].Leaves[0].Implementations[0] != "project.Head.record" { + t.Fatal("nested clone aliases original") + } +} diff --git a/docs/checks/PS6136.md b/docs/checks/PS6136.md new file mode 100644 index 00000000..48fc5a77 --- /dev/null +++ b/docs/checks/PS6136.md @@ -0,0 +1,44 @@ +# PS6136 — maximum-context-retained-output-workspace + + + +**maximum-context output scratch remains resident for a reviewed one-row inference policy** + +- Category: `alloc` +- Fix level: **L3** +- Auto-fix: no (advisory) +- Domain check: requires vocabulary `outputWorkspaceContracts` in perfscan.yaml + +PS6136 requires an exact typed outputWorkspaceContracts entry and proves a selected public model constructor's source flow into a fresh owner, context-times-width float32 output allocation, exact owning backend factory/retention list, current final error barrier, source head Shape width, complete loaded-package constructor/field-use inventory, and selected-instance producer/consumer extents. It does not identify models, projectors or allocators from names. + +The selected common methods must consume one logical output row; an explicit rare bulk method may consume actual token rows. Whole-capacity ToHost is inventoried separately as a physical transfer, with a source-proved first-width logical prefix; it is never counted as a one-row physical leaf. Exact typed selected native buffer/recorder factories, wrappers, forwarding and capability availability are independently bound. Unknown/unstable arguments, aliases, unclassified observations, allocator/geometry changes, unsafe constructor error publication, full logical consumption, contradictory contracts and hard-real-time no-growth policy remain silent. + +The contract supplies the narrowly reviewed model Config width/head relationship, native allocation/access/release/completion semantics, positive checked shapes, provider/build partitions and external observers, sequential ownership and dominant one-row versus rare bulk policy. These are not proved by Go syntax, a source hash or a method name. Unselected architecture/model factories remain outside the selected model-width contract. + +There is no automatic fix. Review one resident output row plus a privately built exact high-water bulk allocation. Preserve all public hot/bulk paths, physical transfer and synchronization, invalid shapes, numeric/reference and sequential/batched parity, partial failure cleanup, final release and publication ownership. A hard-real-time allocator policy may deliberately retain maximum capacity. Measure retained bytes and public constructor/Step/StepNLast/StepN campaigns in order-alternating same-binary runs; do not attribute unrelated projector/backend work to storage alone. + +## Example + +Before: + +```go +d.logits = mk(make([]float32, d.maxLen*d.v)) +// Common logical output: d.v; explicit rare bulk: rows*d.v. +``` + +After: + +```go +// Keep one resident row; privately allocate/reuse a checked exact bulk generation; +// publish only after success; preserve all reads, failures, completion and release. +``` + +## Measured + +Owner issue #887 / GoAI PR #1207 report GPT-2 small context 1024, vocabulary 50257: 205852672 retained float32 output bytes versus 201028 one-row bytes, 205651644 idle bytes and 1024x amplification. Apple M2 Pro same-binary constructor medians were 3423267 versus 157542 ns (21.73x), with roughly 207.9 MB versus 2.25 MB allocations/op. Step median eager/lazy ratio was 1.0067 and StepNLast pp16 1.0115 with unchanged allocation counts; full StepN GPT/Llama parity passed. These are attributed owner measurements, not universal throughput claims or a ranking. + +## Suppressing + +```go +//perfscan:ignore PS6136 +``` diff --git a/docs/checks/README.md b/docs/checks/README.md index 190e1b5b..b2ee4085 100644 --- a/docs/checks/README.md +++ b/docs/checks/README.md @@ -9,7 +9,7 @@ Every check has a stable PS-prefixed ID and a fix level: reporting and fixing: `perfscan -fix` applies exactly the reported checks' fixes. -**429 checks total** — **265 with a bit-identical auto-fix**, 164 advisory. By fix level: **246 L1** (idiomatic), **124 L2** (structured), **59 L3** (aggressive). +**430 checks total** — **265 with a bit-identical auto-fix**, 165 advisory. By fix level: **246 L1** (idiomatic), **124 L2** (structured), **60 L3** (aggressive). | ID | Category | Level | Auto-fix | Title | |----|----------|-------|----------|-------| @@ -439,6 +439,7 @@ checks' fixes. | [PS6132](PS6132.md) | alloc | L3 | | constructor retains maximum-row transient storage despite explicit active-row execution | | [PS6133](PS6133.md) | verify | L3 | | a row-local band offset inflates the whole backing-row shape guard | | [PS6135](PS6135.md) | verify | L3 | | an explicit active-bound capability falls back to capacity-wide work | +| [PS6136](PS6136.md) | alloc | L3 | | maximum-context output scratch remains resident for a reviewed one-row inference policy | | [PS6137](PS6137.md) | alloc | L3 | | native snapshot acquisition freshly materializes owned output | | [PS6138](PS6138.md) | verify | L3 | | an ARM64 counted loop serially extracts adjacent fields with destructive shifts | | [PS7001](PS7001.md) | offload | L3 | | a serial-K GPU reduction kernel (one thread per row, no SIMD-group reduction) leaves lanes idle at M=1 | diff --git a/docs/issue-887-output-workspace.md b/docs/issue-887-output-workspace.md new file mode 100644 index 00000000..ab529cc2 --- /dev/null +++ b/docs/issue-887-output-workspace.md @@ -0,0 +1,174 @@ +# Issue #887: output workspace source audit + +This work targets [perfscan issue #887](https://github.com/jxsl13/perfscan/issues/887). +PS6136 is a registered advisory requiring `outputWorkspaceContracts`. The +existing unregistered PS6125 components remain analysis foundations, not a +replacement for its complete selected-instance proof. No automatic +lifetime-changing rewrite is justified. + +## Authentic owner baseline + +[GoAI PR #1207](https://github.com/jxsl13/goai/pull/1207) changed constructor +output residency in both GPTDecoder and the shared Decoder. The before revision +is `a10a6bff8f7cb0adf695742b6ec677b750b03c28`; the merged after revision is +`40bf79ae7d93caa9384719c779d38b6ae4aaa8a1`. Complete unchanged `gpt.go` and +`decoder.go` files from both revisions are pinned under +`checks/testdata/ps6136-owner/`. Text artifacts are provenance, not evidence of +native execution or successful construction. + +The before GPT constructor stores `mk(make([]float32, c*d.v))` in `d.logits`, +where `c := d.maxLen`. Its captured factory passes the same host slice to +`ops.newBuffer`, appends each successful buffer to this decoder's `all`, and +returns its buffer in a fresh slot. The failure path calls this decoder's +Release. Shared Decoder uses the same allocation expression in allocScratch, +with mkBuf returning its error-tracking, append-to-all factory. Its recurrent +scratch constructors already allocate one vocabulary row and are not maximum +context output candidates. + +## Complete output use obligations + +GPT Step projects one row through its actual head field and downloads `d.v` +float32 values. StepN and StepNLast pass different constant lastOnly flags to +gptStepN. That helper sets rows to len(tokens), changes rows to one on the +last-only branch, and uses that exact rows binding for the head projection and +the `rows*d.v` download. Generate uses StepNLast followed by Step. + +Shared Decoder additionally routes projection through recordLogits, optionally +adds bias over `rows*d.v`, and pre-encodes single-row work in encodeStep. +The Metal ProfileMetalStep wrapper calls ordinary Step on the same decoder, +temporarily changes only recorder/async members, and restores them afterward. +Its physical projection still uses the original mRec MatMul/AddBias forwarding +operations; the buffer allocator is not replaced. It is a relevant one-row API, +not an omitted alias or an independent decoder. +stepN's recurrent branch uses sequential Step calls rather than a multirow +device output. The batched branch specializes rows using lastOnly and uses +that same row extent for projection and download. + +Generate's device capability paths call TopKN and SoftmaxStatsN with `d.v`. +Its pure-top-p fallback calls **ToHost with no active extent**, then copies only +the first `d.v` values into the sampler input. This must be inventoried as a +capacity-wide physical transfer with one-row logical consumption. It is not +evidence that every physical backend operation is bounded to one row. PR1207 +preserves that fallback after reducing common resident storage. + +An API audit of all twelve other non-test Go files in the pinned llamagpu +package found additional `.logits` accesses only in cuda_graph_llama.go. Those +belong to a distinct graph owner and its constructor already requests +cfg.Vocab elements. Typed ownership identity must exclude them; spelling alone +is not a match. Analysis must still close observations in every loaded package +file and account for other build paths through explicit reviewed facts. +The mandatory Metal owner replay includes all 28 shared and two GPT constructor +functions (including the unselected quantized wrappers), plus the profiling API; +only selected GPT and Llama model/head/provider relationships justify findings. + +## Source versus reviewed semantics + +Source must establish the actual constructor/factory/retained-buffer flow, +typed field identities, projector factory flow, every producer and consumer +argument, specialized row geometry, width identity, lifecycle call flow and +absence of contradictory uses. Opaque native allocation, projection width +versus model Vocab, active-prefix backend semantics, synchronization and +external/reflection observations require narrowly scoped reviewed contracts. +Workload common-versus-bulk policy is also reviewed, not inferred from names, +method visibility or static call frequency. A full logical consumer or a +hard-real-time prohibition on growth allocation invalidates this advice. + +## Residency and validation + +For float32 output, symbolic retained bytes are `4*Ctx*Vocab`, common bytes +`4*Vocab`, avoidable idle bytes `4*(Ctx-1)*Vocab`, and amplification `Ctx` for +one-row common work. Checked numeric evaluation of the owner's GPT-2-small +profile (Ctx 1024, Vocab 50257) gives 205852672 retained, 201028 common and +205651644 idle bytes, with 1024x amplification. These are an attributed model +profile, not values inferred for every decoder or a ranking claim. + +The owner reports same-binary M2 Pro constructor medians 3423267 versus 157542 +ns/op (21.73x), with allocation measurements approximately 207905393–207906016 +versus 2253424–2254015 B/op. Order-alternated public Step and StepNLast median +ratios were 1.0067 and 1.0115, with full-StepN GPT and Llama parity passing. +Those are project measurements, not universal throughput guarantees. + +A remedy retains one common row and lazily grows a separate exact reusable +bulk high-water buffer. Preserve Step, StepN and StepNLast outputs, pending +recorders and completion, all optional providers and fallbacks, checked size +arithmetic, invalid shapes, failed growth cleanup, replacement and final +release. Measure constructor residency and latency against a same-binary eager +control, plus order-alternated complete public decode/prefill latency and +allocations. Native parity and lifecycle tests are necessary; source fixtures +do not execute kernels or replace those validation gates. + +## Registered source closure + +The registered replay compiles the unchanged before/after owner files and +selected public Metal factories. Both original GPT and Llama constructors +produce one diagnostic; both actual PR1207 after constructors produce none. +Selected public NewGPT/New model arguments flow to the exact private factory, +Config.Ctx/Vocab flow to the same fresh owner's maxLen/v, and the actual +Head/Out Shape()[1] flows to f32Linear.n. The selected New backendOps literal +omits quantizeF32: typed closed zero-field specialization proves the F32 branch, +not an invented width for native quantized weights. Other model families are +outside this selected width contract and stay silent. + +The source proves the exact NewDeviceBufferF32 host descriptor, returned mBuf +type/native field and mb unwrapping, plus newRecorder -> NewRecorder -> mRec.r. +f32Linear.record preserves output, active rows and stored dimensions to the +configured recorder method; mRec forwards every buffer/geometry role to the +exact native API. ProfileMetalStep binds NewProfilingRecorder(maxEvents) to the +same embedded mRec.r and separate local metadata. Its exact deferred frame +restores both callback cells and async state. Unknown override/restore or late +work is rejected. All relevant adapter methods resolve to the same base method +objects, not merely a second field of the same pointer type. + +Metal's exact mBuf type implements deviceTopKer but not deviceTopPer. A typed +capability-absence witness may exclude the latter from *selected-provider +execution*, never from the full source inventory. Its latent ToHost operation +still has explicitly reviewed capacity-wide physical semantics and a +source-proved first-Vocab logical copy. No Metal ToHost method is fabricated, +and no observed common physical transfer ratio is claimed. + +Complete owner constructor/field/whole-owner exposure closure is unconditional +over the loaded package. Source-call contexts cannot share an approved +leaf binding when one invocation has unknown rows/projector flow. Native +callback/list mutations are checked per published invocation too: a helper's +safe fresh-constructor use does not justify its use on a published owner. +Each actual output recorder receiver (or projector recorder argument) must +originate from this same owner's primary/secondary recorder factory result; +every reachable factory phi alternative is checked as a raw owner-field read. +A matching interface signature or a separately reviewed adapter is insufficient. +The rare bulk path must contain a source-bound actualRows*width extent witness; +policy configuration cannot invent a full-row requirement for a one-row path. +Constructor workspace reads and unclassified retention-list aliases reject; +the list permits only canonical append-base reads and exact Release loop/clear +with no control-flow prefix capable of bypassing cleanup. +The final checked error load/comparison/branch must be current and effect-free, +with no post-guard allocation/error write before publication. + +The conventional analysistest fixture retains the complete original GPT owner +file and replays the full selected-package observation inventory with explicit +type-only model/native scaffolds. Required native method receivers are checked +after typed resolution: a dotted package path cannot disguise a package-level +function as a method. Regression controls reproduced the receiver panic and +both owners' early-return cleanup false positives before their repairs. +Additional independent controls reproduced an effectful copy destination that +mutated its iterator, arbitrary-recorder output leaves and a one-row-only bulk +path. The canonical logical copy now permits exactly its four read/increment +iterator uses and a bare typed destination, with no iterator escape or effectful +destination evaluation. Six registered altered-owner controls remain alongside +the genuine before/after positives. + +Contracts attest only the reviewed opaque and policy facts: model width/head +relationship; fresh exclusive F32 native ownership; exact native access, +release, completion and synchronization; checked shape/index conversion; +sequential lifecycle; audited provider/build partitions and external/unsafe/ +reflection observers; dominant one-row and explicit rare bulk policy. Exact +typed identity, source buffer/row/head/call flow and full loaded observations +are mandatory despite those facts. Hard-real-time no-growth policy or full +logical consumers suppress the candidate. Symbolic metrics are conditional on +positive checked width and maximumRows > 1; numeric profiles are attributed +configuration, not runtime dimensions inferred by the analyzer. + +The shared-helper regression was also run against a temporary mechanical +pre-repair approval-map overlay: it failed with an unsafe candidate (checks +1.080s). Context-keyed maps passed the retained regression and compound owner +tests (3.231s). The overlay changed only approval/read keys, not fixture tests +or source semantics. diff --git a/internal/runner/ps6136_vocabulary_test.go b/internal/runner/ps6136_vocabulary_test.go new file mode 100644 index 00000000..2a78ed1a --- /dev/null +++ b/internal/runner/ps6136_vocabulary_test.go @@ -0,0 +1,16 @@ +package runner + +import ( + "github.com/jxsl13/perfscan/checks" + "github.com/jxsl13/perfscan/config" + "testing" +) + +func TestPS6136StarvedVocabulary(t *testing.T) { + t.Parallel() + for _, cfg := range []config.Config{{}, {OutputWorkspaceContracts: []config.OutputWorkspaceContract{{OwnerType: "project.Owner", WorkspaceField: "output"}}}} { + if missing := missingVocab(checks.PS6136, &cfg); len(missing) != 1 || missing[0] != "outputWorkspaceContracts" { + t.Fatalf("missing %v", missing) + } + } +} diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 754a6269..541c8526 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -444,6 +444,7 @@ func missingVocab(c *lint.Check, cfg *config.Config) []string { "inPlaceFusionContracts": len(cfg.InPlaceFusionContracts), "receiverStagingContracts": validReceiverStagingContracts(cfg.ReceiverStagingContracts), "activeBoundFallbackContracts": config.UsableActiveBoundFallbackContractCount(cfg.ActiveBoundFallbackContracts), + "outputWorkspaceContracts": config.UsableOutputWorkspaceContractCount(cfg.OutputWorkspaceContracts), "nativeSnapshotReuseContracts": config.UsableNativeSnapshotReuseContractCount(cfg.NativeSnapshotReuseContracts), "reusableResultLoopContracts": validReusableResultLoopContracts(cfg.ReusableResultLoopContracts), "recorderResidualAddContracts": validRecorderResidualAddContracts(cfg.RecorderResidualAddContracts), From 0f2ff90b499b0ed643a621980200a15ae122ce54 Mon Sep 17 00:00:00 2001 From: John Behm Date: Sat, 12 Sep 2026 19:32:10 +0200 Subject: [PATCH 2/3] fix: validate PS6136 source parameter identities --- checks/ps6136.go | 20 ++++++++++------ checks/ps6136_compiler_test.go | 4 +--- checks/ps6136_parameter_object.go | 17 +++++++++++++ checks/ps6136_parameter_object_test.go | 33 ++++++++++++++++++++++++++ 4 files changed, 64 insertions(+), 10 deletions(-) create mode 100644 checks/ps6136_parameter_object.go create mode 100644 checks/ps6136_parameter_object_test.go diff --git a/checks/ps6136.go b/checks/ps6136.go index fd2e3c62..4f1ec357 100644 --- a/checks/ps6136.go +++ b/checks/ps6136.go @@ -266,19 +266,24 @@ func (selection *ps6136Selection) consumerCalls(context *ps6125SSAContext, leave return nil } parameter := context.flow.function.Params[0] - if !types.Identical(parameter.Type(), types.NewPointer(selection.ownerType)) || parameter.Object() == nil { + parameterObject := ps6136ParameterObject(parameter) + if parameterObject == nil || !types.Identical(parameter.Type(), types.NewPointer(selection.ownerType)) { return nil } owner := context.reference(parameter) var pool ps6125ExtentPool - width := ps6125SymbolicExtent(pool.identity(parameter.Object(), []*types.Var{selection.width}, false)) + width := ps6125SymbolicExtent(pool.identity(parameterObject, []*types.Var{selection.width}, false)) var bulkExtent ps6125Extent if bulk { - if len(context.flow.function.Params) < 2 || !types.Identical(context.flow.function.Params[1].Type(), types.NewSlice(types.Typ[types.Int])) { + if len(context.flow.function.Params) < 2 { return nil } tokens := context.flow.function.Params[1] - length := ps6125SymbolicExtent(pool.identity(tokens.Object(), nil, true)) + tokensObject := ps6136ParameterObject(tokens) + if tokensObject == nil || !types.Identical(tokens.Type(), types.NewSlice(types.Typ[types.Int])) { + return nil + } + length := ps6125SymbolicExtent(pool.identity(tokensObject, nil, true)) // Recompute the contextual scalar/branch algebra after binding the // actual bulk descriptor, retaining the same invocation identity. context.flow = ps6125AnalyzeSSAExtents(context.flow.function, nil, map[*ssa.Parameter]ps6125Extent{tokens: length}) @@ -498,12 +503,13 @@ func (selection *ps6136Selection) constructorAllocation(remaining int) *ps6136Co return nil } model, ok := selection.model.value.(*ssa.Parameter) - if !ok || model.Object() == nil { + modelObject := ps6136ParameterObject(model) + if !ok || modelObject == nil { return nil } var pool ps6125ExtentPool - rows := pool.identity(model.Object(), []*types.Var{selection.modelConfig, selection.configRows}, false) - width := pool.identity(model.Object(), []*types.Var{selection.modelConfig, selection.configWidth}, false) + rows := pool.identity(modelObject, []*types.Var{selection.modelConfig, selection.configRows}, false) + width := pool.identity(modelObject, []*types.Var{selection.modelConfig, selection.configWidth}, false) residency, ok := ps6136OutputResidency(rows, width) if !ok { return nil diff --git a/checks/ps6136_compiler_test.go b/checks/ps6136_compiler_test.go index 78a27edc..598bfaca 100644 --- a/checks/ps6136_compiler_test.go +++ b/checks/ps6136_compiler_test.go @@ -81,9 +81,7 @@ func sampleTopKCandidates(*nlp.Sampler,[]int32,[]float32)int{return 0} return nil, err } // Owner bodies keep their original import binding environment. - for _, declaration := range extra.Decls { - support.Decls = append(support.Decls, declaration) - } + support.Decls = append(support.Decls, extra.Decls...) support.Imports = append(support.Imports, extra.Imports...) deviceSource, err := os.ReadFile("testdata/ps6136-owner/sampling_fastpath.go.txt") if err != nil { diff --git a/checks/ps6136_parameter_object.go b/checks/ps6136_parameter_object.go new file mode 100644 index 00000000..27104f72 --- /dev/null +++ b/checks/ps6136_parameter_object.go @@ -0,0 +1,17 @@ +package checks + +import ( + "go/types" + + "golang.org/x/tools/go/ssa" +) + +// Parameter.Object wraps its source *types.Var in an interface, including a +// typed nil for zero/incomplete parameters. Extent roots require a real source var. +func ps6136ParameterObject(parameter *ssa.Parameter) *types.Var { + if parameter == nil { + return nil + } + object, _ := parameter.Object().(*types.Var) + return object +} diff --git a/checks/ps6136_parameter_object_test.go b/checks/ps6136_parameter_object_test.go new file mode 100644 index 00000000..2b9ce8a5 --- /dev/null +++ b/checks/ps6136_parameter_object_test.go @@ -0,0 +1,33 @@ +package checks + +import ( + "testing" + + "golang.org/x/tools/go/ssa" +) + +func TestPS6136ParameterObject(t *testing.T) { + t.Parallel() + for _, parameter := range []*ssa.Parameter{nil, {}} { + if got := ps6136ParameterObject(parameter); got != nil { + t.Fatalf("synthetic parameter returned source identity: %v", got) + } + } + fixture, err := ps6136CompileOwners(t, "before", true, true) + if err != nil { + t.Fatal(err) + } + parameter := ps6136FixtureSSA(fixture).Func("newGPTDecoder").Params[0] + got := ps6136ParameterObject(parameter) + if got == nil || got != parameter.Object() { + t.Fatalf("lost exact source parameter identity: %v", got) + } +} + +func TestPS6136NilConsumerParameterRejects(t *testing.T) { + t.Parallel() + context := &ps6125SSAContext{flow: &ps6125SSAExtents{function: &ssa.Function{Params: []*ssa.Parameter{nil}}}} + if got := (&ps6136Selection{}).consumerCalls(context, nil, false, nil); got != nil { + t.Fatal("nil receiver parameter did not reject") + } +} From 422b759ce0869ae786cea1ed73e6024a28a8299d Mon Sep 17 00:00:00 2001 From: John Behm Date: Sat, 12 Sep 2026 20:55:43 +0200 Subject: [PATCH 3/3] fix(PS6136): prove constructor geometry read ordering --- checks/ps6136.go | 28 +++- checks/ps6136_compiler_test.go | 14 ++ checks/ps6136_initialization_context_test.go | 65 ++++++++++ checks/ps6136_initialization_order.go | 48 +++++++ ...ps6136_initialization_order_review_test.go | 122 ++++++++++++++++++ docs/issue-887-output-workspace.md | 6 + 6 files changed, 282 insertions(+), 1 deletion(-) create mode 100644 checks/ps6136_initialization_context_test.go create mode 100644 checks/ps6136_initialization_order.go create mode 100644 checks/ps6136_initialization_order_review_test.go diff --git a/checks/ps6136.go b/checks/ps6136.go index 4f1ec357..78b40cd3 100644 --- a/checks/ps6136.go +++ b/checks/ps6136.go @@ -447,6 +447,13 @@ func (selection *ps6136Selection) initializationValues(remaining int) bool { return false } counts := make(map[*types.Var]int) + type observation struct { + context *ps6125SSAContext + instruction ssa.Instruction + field *types.Var + } + initializers := make(map[*types.Var]observation) + var reads []observation pending := []*ps6125SSAContext{selection.context} for len(pending) != 0 { context := pending[len(pending)-1] @@ -461,6 +468,15 @@ func (selection *ps6136Selection) initializationValues(remaining int) bool { continue } for _, instruction := range block.Instrs { + if load, ok := instruction.(*ssa.UnOp); ok && load.Op == token.MUL { + path := paths.resolve(load.X) + if path.known && len(path.access.fields) == 1 && ps6136AccessOwnerRoot(context, path, selection.ownerType) == selection.owner { + field := path.access.fields[0] + if field == selection.maximumRows || field == selection.width { + reads = append(reads, observation{context, load, field}) + } + } + } if store, ok := instruction.(*ssa.Store); ok { path := paths.resolve(store.Addr) if path.known && len(path.access.fields) == 1 && ps6136AccessOwnerRoot(context, path, selection.ownerType) == selection.owner { @@ -474,6 +490,7 @@ func (selection *ps6136Selection) initializationValues(remaining int) bool { return false } counts[field]++ + initializers[field] = observation{context, store, field} } } } @@ -488,7 +505,16 @@ func (selection *ps6136Selection) initializationValues(remaining int) bool { } } } - return counts[selection.maximumRows] == 1 && counts[selection.width] == 1 + if counts[selection.maximumRows] != 1 || counts[selection.width] != 1 { + return false + } + for _, read := range reads { + initializer := initializers[read.field] + if !ps6136ContextInstructionDominates(initializer.context, initializer.instruction, read.context, read.instruction, 128) { + return false + } + } + return true } type ps6136ConstructorProof struct { diff --git a/checks/ps6136_compiler_test.go b/checks/ps6136_compiler_test.go index 598bfaca..314256e6 100644 --- a/checks/ps6136_compiler_test.go +++ b/checks/ps6136_compiler_test.go @@ -230,6 +230,20 @@ func sampleTopKCandidates(*nlp.Sampler,[]int32,[]float32)int{return 0} } if rewrite != nil { files = rewrite(fs, files) + // Meaningful injected nodes must have real source positions: a NoPos + // call can otherwise fail source census and vacuously pass a negative. + for index, file := range files { + var source bytes.Buffer + if err := format.Node(&source, fs, file); err != nil { + return nil, err + } + name := fs.Position(file.Pos()).Filename + parsed, err := parser.ParseFile(fs, name, source.Bytes(), parser.ParseComments|parser.SkipObjectResolution) + if err != nil { + return nil, err + } + files[index] = parsed + } } pkg, err := (&types.Config{Importer: imports}).Check("github.com/jxsl13/goai/llamagpu", fs, files, info) return &analysisOwnerFixture{fs, files, info, pkg, imports.sources}, err diff --git a/checks/ps6136_initialization_context_test.go b/checks/ps6136_initialization_context_test.go new file mode 100644 index 00000000..05d6e429 --- /dev/null +++ b/checks/ps6136_initialization_context_test.go @@ -0,0 +1,65 @@ +package checks + +import ( + "go/token" + "testing" + + "golang.org/x/tools/go/ssa" +) + +func TestPS6136InitializationInvocationOrdering(t *testing.T) { + t.Parallel() + pkg := ps6125TestSSA(t, `package extent +func leaf(int){} +func initialize(p *int){*p=1} +func conditional(p *int,b bool){if b {*p=1}} +func read(p *int){leaf(*p)} +func callback(fn func(*int),p *int){fn(p)} +func early(){p:=new(int);initialize(p);read(p)} +func late(){p:=new(int);read(p);initialize(p)} +func guarded(b bool){p:=new(int);conditional(p,b);read(p)} +func callbackEarly(){p:=new(int);callback(initialize,p);read(p)} +func callbackLate(){p:=new(int);read(p);callback(initialize,p)} +func aliasLate(){p:=new(int);q:=p;read(q);initialize(p)} +`) + for _, name := range []string{"early", "late", "guarded", "callbackEarly", "callbackLate", "aliasLate"} { + t.Run(name, func(t *testing.T) { + t.Parallel() + root := ps6125NewSSAContext(pkg.Func(name), nil, nil, 128) + var store *ssa.Store + var load *ssa.UnOp + var storeContext, loadContext *ps6125SSAContext + pending := []*ps6125SSAContext{root} + for len(pending) > 0 { + context := pending[len(pending)-1] + pending = pending[:len(pending)-1] + for _, block := range context.flow.function.Blocks { + if !context.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + switch instruction := instruction.(type) { + case *ssa.Store: + store, storeContext = instruction, context + case *ssa.UnOp: + if instruction.Op == token.MUL { + load, loadContext = instruction, context + } + case *ssa.Call: + if child := ps6136Call(context, instruction); child != nil { + pending = append(pending, child) + } + } + } + } + } + if store == nil || load == nil { + t.Fatal("typed store/load boundary missing") + } + wanted := name == "early" || name == "callbackEarly" + if got := ps6136ContextInstructionDominates(storeContext, store, loadContext, load, 128); got != wanted { + t.Fatalf("ordered=%v want %v", got, wanted) + } + }) + } +} diff --git a/checks/ps6136_initialization_order.go b/checks/ps6136_initialization_order.go new file mode 100644 index 00000000..a09cc339 --- /dev/null +++ b/checks/ps6136_initialization_order.go @@ -0,0 +1,48 @@ +package checks + +import "golang.org/x/tools/go/ssa" + +// ps6136ContextInstructionDominates orders exact source invocation contexts. +// A callee write is lifted to its call only if it dominates every executable +// return; an ancestor write must dominate the actual descendant invocation. +// Unknown completion paths are not an initialization guarantee. +func ps6136ContextInstructionDominates(beforeContext *ps6125SSAContext, before ssa.Instruction, afterContext *ps6125SSAContext, after ssa.Instruction, remaining int) bool { + if beforeContext == nil || afterContext == nil || beforeContext.flow == nil || afterContext.flow == nil || before == nil || after == nil || before.Parent() != beforeContext.flow.function || after.Parent() != afterContext.flow.function || !beforeContext.flow.blocks[before.Block()] || !afterContext.flow.blocks[after.Block()] || remaining <= 0 { + return false + } + points := make(map[*ps6125SSAContext]ssa.Instruction) + for current, point := afterContext, after; current != nil; current, point = current.parent, current.site { + remaining-- + if remaining <= 0 { + return false + } + points[current] = point + } + for current, point := beforeContext, before; current != nil; current, point = current.parent, current.site { + remaining-- + if remaining <= 0 { + return false + } + if target := points[current]; target != nil { + return current.flow.instructionDominates(point, target) + } + returned := false + for _, block := range current.flow.function.Blocks { + if !current.flow.blocks[block] { + continue + } + for _, instruction := range block.Instrs { + if _, ok := instruction.(*ssa.Return); ok { + if !current.flow.instructionDominates(point, instruction) { + return false + } + returned = true + } + } + } + if !returned || current.site == nil { + return false + } + } + return false +} diff --git a/checks/ps6136_initialization_order_review_test.go b/checks/ps6136_initialization_order_review_test.go new file mode 100644 index 00000000..aa225add --- /dev/null +++ b/checks/ps6136_initialization_order_review_test.go @@ -0,0 +1,122 @@ +package checks + +import ( + "bytes" + "go/ast" + "go/parser" + "go/printer" + "go/token" + "testing" + + "github.com/jxsl13/perfscan/config" + "golang.org/x/tools/go/analysis" +) + +// Pre-repair, all four fully reprinted/reparsed authentic controls emitted one +// registered diagnostic (checks 1.872s). Direct NoPos AST injection originally +// appeared to reject; real-source normalization exposed that vacuous result. +func TestPS6136RegisteredInitializationOrderReview(t *testing.T) { + t.Parallel() + for _, field := range []string{"maxLen", "v"} { + for _, conditional := range []bool{false, true} { + t.Run(field+map[bool]string{false: "/late", true: "/conditional"}[conditional], func(t *testing.T) { + t.Parallel() + fixture, err := ps6136CompileOwnersRewrite(t, "before", true, true, func(fs *token.FileSet, files []*ast.File) []*ast.File { + var value ast.Expr + var constructor *ast.FuncDecl + for _, file := range files { + for _, decl := range file.Decls { + fn, ok := decl.(*ast.FuncDecl) + if !ok || fn.Name.Name != "newGPTDecoder" { + continue + } + constructor = fn + ast.Inspect(fn.Body, func(n ast.Node) bool { + lit, ok := n.(*ast.CompositeLit) + if !ok { + return true + } + id, ok := lit.Type.(*ast.Ident) + if !ok || id.Name != "GPTDecoder" { + return true + } + for i, elt := range lit.Elts { + kv, ok := elt.(*ast.KeyValueExpr) + if ok && kv.Key.(*ast.Ident).Name == field { + value = kv.Value + lit.Elts = append(lit.Elts[:i], lit.Elts[i+1:]...) + break + } + } + return true + }) + } + } + if constructor == nil || value == nil { + t.Fatal("exact GPT initializer missing") + } + assignment := &ast.AssignStmt{Lhs: []ast.Expr{&ast.SelectorExpr{X: ast.NewIdent("d"), Sel: ast.NewIdent(field)}}, Tok: token.ASSIGN, Rhs: []ast.Expr{value}} + var inserted ast.Stmt = assignment + if conditional { + inserted = &ast.IfStmt{Cond: &ast.BinaryExpr{X: &ast.SelectorExpr{X: ast.NewIdent("cfg"), Sel: ast.NewIdent("Ctx")}, Op: token.GTR, Y: &ast.BasicLit{Kind: token.INT, Value: "1"}}, Body: &ast.BlockStmt{List: []ast.Stmt{assignment}}} + } + count := 0 + for i, statement := range constructor.Body.List { + assign, ok := statement.(*ast.AssignStmt) + if !ok || len(assign.Lhs) != 1 { + continue + } + selector, ok := assign.Lhs[0].(*ast.SelectorExpr) + if !ok || selector.Sel.Name != "logits" { + continue + } + constructor.Body.List = append(constructor.Body.List[:i+1], append([]ast.Stmt{inserted}, constructor.Body.List[i+1:]...)...) + count++ + break + } + if count != 1 { + t.Fatal("exact logits allocation missing") + } + for i, file := range files { + for _, decl := range file.Decls { + if decl != constructor { + continue + } + var source bytes.Buffer + if err := printer.Fprint(&source, fs, file); err != nil { + t.Fatal(err) + } + parsed, err := parser.ParseFile(fs, "order-review-gpt.go", source.Bytes(), parser.SkipObjectResolution) + if err != nil { + t.Fatal(err) + } + files[i] = parsed + } + } + return files + }) + if err != nil { + t.Fatal(err) + } + pkg := ps6136FixtureSSA(fixture) + contract := ps6136CompleteOwnerContract(pkg, "GPTDecoder") + var diagnostics []analysis.Diagnostic + pass := &analysis.Pass{Fset: fixture.fileset, Files: fixture.files, TypesInfo: fixture.info, Pkg: fixture.pkg, Report: func(d analysis.Diagnostic) { diagnostics = append(diagnostics, d) }} + context := ps6136ContractsContext(pass) + selection := context.selection(pkg, context.callable(contract.ConstructorEntries[0]), &contract) + if selection == nil { + t.Fatal("selected genuine constructor missing") + } + if selection.initializationValues(16384) { + t.Fatal("late or conditional geometry accepted at initialization gate") + } + if _, err := runPS6136WithContracts(pass, []config.OutputWorkspaceContract{contract}); err != nil { + t.Fatal(err) + } + if len(diagnostics) != 0 { + t.Fatalf("late or conditional initialization produced %d findings", len(diagnostics)) + } + }) + } + } +} diff --git a/docs/issue-887-output-workspace.md b/docs/issue-887-output-workspace.md index ab529cc2..f9109c95 100644 --- a/docs/issue-887-output-workspace.md +++ b/docs/issue-887-output-workspace.md @@ -73,6 +73,12 @@ Workload common-versus-bulk policy is also reviewed, not inferred from names, method visibility or static call frequency. A full logical consumer or a hard-real-time prohibition on growth allocation invalidates this advice. +The unique source initialization of maximum rows and width must precede every +selected constructor geometry load. Exact helper invocation boundaries are +ordered too: a helper write is complete only when it dominates all executable +returns. Late assignments and conditional initialization cannot establish the +allocation extent, even if the eventual published owner has correct fields. + ## Residency and validation For float32 output, symbolic retained bytes are `4*Ctx*Vocab`, common bytes