{ "$schema": "./kernel-registry.schema.json", "version": "1.0", "operations": { "matmul": { "description": "Matrix multiplication kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "A", "type": "read-only-storage" }, { "index": 2, "name": "B", "type": "read-only-storage" }, { "index": 3, "name": "C", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "M", "type": "u32", "offset": 0 }, { "name": "N", "type": "u32", "offset": 4 }, { "name": "K", "type": "u32", "offset": 8 }, { "name": "alpha", "type": "f32", "offset": 12 } ] }, "variants": { "f16": { "wgsl": "matmul_f16.wgsl", "entryPoint": "main", "workgroup": [16,16,1], "requires": ["shader-f16"], "outputDtype": "f16", "weightDtype": "f16", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#matmulKernel"], "wgslOnDisk": true } }, "f16_vec4": { "wgsl": "matmul_f16.wgsl", "entryPoint": "main_vec4", "workgroup": [16,16,1], "requires": ["shader-f16"], "outputDtype": "f16", "weightDtype": "f16", "variantMetadata": { "colsPerThread": 4 }, "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#matmulKernel"], "wgslOnDisk": true } }, "f16_tiled": { "wgsl": "matmul_f16_tiled.wgsl", "entryPoint": "main", "workgroup": [16,16,1], "requires": ["shader-f16"], "outputDtype": "f16", "weightDtype": "f16", "variantMetadata": { "tileM": 64, "tileN": 64 }, "description": "Register-tiled prefill matmul (transpose_b=1 row-major B): 4x4 thread tile, 64x64 workgroup tile, 16x outputs/workgroup vs base", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#matmulKernel"], "wgslOnDisk": true } }, "f16w_f32a": { "wgsl": "matmul_f16w_f32a.wgsl", "entryPoint": "main", "workgroup": [16,16,1], "requires": ["shader-f16"], "outputDtype": "f32", "weightDtype": "f16", "reachability": { "status": "pinned", "inlineConfigs": [ "amplify/amplify-120m-f16-af32", "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "embeddinggemma/google-embeddinggemma-300m-q4k-ehf16-af32", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "minicpm/minicpm4-0-5b-f16-af32", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": ["matmul.rules.json#matmulKernel"], "wgslOnDisk": true } }, "f16w_f32a_tiled": { "wgsl": "matmul_f16w_f32a_tiled.wgsl", "entryPoint": "main", "workgroup": [16,16,1], "requires": ["shader-f16"], "outputDtype": "f32", "weightDtype": "f16", "variantMetadata": { "tileM": 64, "tileN": 64 }, "description": "Register-tiled prefill matmul: f32 activations, f16 weights, f32 output. 4x4 thread tile, 64x64 workgroup tile.", "reachability": { "status": "pinned", "inlineConfigs": [ "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32" ], "ruleChains": ["matmul.rules.json#matmulKernel"], "wgslOnDisk": true } }, "gemv": { "wgsl": "matmul_gemv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "weightDtype": "f16", "description": "Optimized GEMV for M=1 decode: uses shared memory for A vector", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#gemvVariant"], "wgslOnDisk": true } }, "gemv_f16a": { "wgsl": "matmul_gemv_f16a.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "weightDtype": "f16", "description": "Optimized GEMV for M=1 decode with f16 activations + f16 output", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#gemvVariant"], "wgslOnDisk": true } }, "gemv_subgroup": { "wgsl": "matmul_gemv_subgroup.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f32", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 4 }, "description": "Subgroup-optimized GEMV - 1.5x faster using subgroupAdd", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#gemvPathVariant","matmul.rules.json#gemvVariant"], "wgslOnDisk": true } }, "gemv_subgroup_f16a": { "wgsl": "matmul_gemv_subgroup_f16a.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 4 }, "description": "Subgroup-optimized GEMV with f16 activations + f16 output", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#gemvPathVariant","matmul.rules.json#gemvVariant"], "wgslOnDisk": true } }, "gemv_subgroup_vec4": { "wgsl": "matmul_gemv_subgroup.wgsl", "entryPoint": "main_vec4", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f32", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 4 }, "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32", "minicpm/minicpm4-0-5b-f16-af32", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": ["matmul.rules.json#gemvVariant"], "wgslOnDisk": true } }, "gemv_subgroup_vec4_f16a": { "wgsl": "matmul_gemv_subgroup_f16a.wgsl", "entryPoint": "main_vec4", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 4 }, "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#gemvVariant"], "wgslOnDisk": true } }, "gemv_subgroup_cols8_f16a": { "wgsl": "matmul_gemv_subgroup_f16a.wgsl", "entryPoint": "main_cols8", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 8 }, "description": "Subgroup GEMV with 8 columns per workgroup (f16 activations/output)", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "gemv_subgroup_vec4_cols8_f16a": { "wgsl": "matmul_gemv_subgroup_f16a.wgsl", "entryPoint": "main_vec4_cols8", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 8 }, "description": "Subgroup GEMV vec4 with 8 columns per workgroup (f16 activations/output)", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "gemv_subgroup_multicol": { "wgsl": "matmul_gemv_subgroup.wgsl", "entryPoint": "main_multicol", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f32", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 8 }, "description": "Multi-column GEMV for large vocab (LM head F16) - 32 columns per workgroup", "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32", "minicpm/minicpm4-0-5b-f16-af32", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": ["matmul.rules.json#gemvPathVariant","matmul.rules.json#gemvVariant"], "wgslOnDisk": true } }, "gemv_subgroup_multicol_f16a": { "wgsl": "matmul_gemv_subgroup_f16a.wgsl", "entryPoint": "main_multicol", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 8 }, "description": "Multi-column GEMV for large vocab (LM head F16) - f16 activations/output", "reachability": { "status": "pinned", "inlineConfigs": ["qwen3/qwen-3-5-0-8b-q4k-ehaf16"], "ruleChains": ["matmul.rules.json#gemvPathVariant","matmul.rules.json#gemvVariant"], "wgslOnDisk": true } }, "gemv_subgroup_cols64": { "wgsl": "matmul_gemv_subgroup.wgsl", "entryPoint": "main_cols64", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f32", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 64 }, "description": "Cooperative 4-thread GEMV - 64 columns per workgroup, 4 threads per column (experimental)", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "gemv_subgroup_cols64_f16a": { "wgsl": "matmul_gemv_subgroup_f16a.wgsl", "entryPoint": "main_cols64", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 64 }, "description": "Cooperative 4-thread GEMV f16a - 64 columns per workgroup, 4 threads per column (experimental)", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "gemv_subgroup_vec4_cols64_f16a": { "wgsl": "matmul_gemv_subgroup_f16a.wgsl", "entryPoint": "main_vec4_cols64", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "f16", "variantMetadata": { "colsPerWg": 64 }, "description": "Cooperative 4-thread GEMV f16a vec4 - 64 columns per workgroup with vec4 loads (experimental)", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_f16a": { "wgsl": "fused_matmul_q4_f16a.wgsl", "entryPoint": "main_f16a", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "q4k", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "litert_int4_multicol_f16a": { "wgsl": "fused_matmul_litert_int4_f16a.wgsl", "entryPoint": "main_multicol_f16a", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "weightDtype": "litert_int4", "variantMetadata": { "colsPerWg": 8 }, "description": "LiteRT fixed-affine INT4 batched matmul with f16 activations and f16 output; keeps Gemma 4 split-package weights packed on GPU", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#litertInt4FusedVariant"], "wgslOnDisk": true } }, "litert_int4_multicol_f16a_f32out": { "wgsl": "fused_matmul_litert_int4_f16a_f32out.wgsl", "entryPoint": "main_multicol_f16a_f32out", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "weightDtype": "litert_int4", "variantMetadata": { "colsPerWg": 8 }, "description": "LiteRT fixed-affine INT4 batched matmul with f16 activations and f32 output; keeps Gemma 4 split-package weights packed on GPU", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#litertInt4FusedVariant"], "wgslOnDisk": true } }, "litert_int4_multicol_f32a_f32out": { "wgsl": "fused_matmul_litert_int4_f32a_f32out.wgsl", "entryPoint": "main_multicol_f32a_f32out", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "weightDtype": "litert_int4", "variantMetadata": { "colsPerWg": 8 }, "description": "LiteRT fixed-affine INT4 batched matmul with f32 activations and f32 output; keeps Gemma 4 split-package logits packed on GPU", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#litertInt4FusedVariant"], "wgslOnDisk": true } }, "w4a16_multicol_f16a": { "wgsl": "fused_matmul_w4a16_f16a.wgsl", "entryPoint": "main_multicol_f16a", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "weightDtype": "w4a16", "variantMetadata": { "colsPerWg": 8 }, "description": "Compressed-tensors W4A16 packed matmul with f16 activations and f16 output; keeps QAT weights packed on GPU with per-row-group scales", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#w4a16FusedVariant"], "wgslOnDisk": true } }, "w4a16_multicol_f16a_f32out": { "wgsl": "fused_matmul_w4a16_f16a_f32out.wgsl", "entryPoint": "main_multicol_f16a_f32out", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "weightDtype": "w4a16", "variantMetadata": { "colsPerWg": 8 }, "description": "Compressed-tensors W4A16 packed matmul with f16 activations and f32 output; keeps QAT weights packed on GPU with per-row-group scales", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#w4a16FusedVariant"], "wgslOnDisk": true } }, "w4a16_batched_f16a": { "wgsl": "fused_matmul_w4a16_batched_f16a.wgsl", "entryPoint": "main_batched_f16a", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "weightDtype": "w4a16", "variantMetadata": { "colsPerWg": 8, "tileM": 4 }, "description": "Compressed-tensors W4A16 packed batched prefill matmul with f16 activations and f16 output; keeps QAT weights packed on GPU with per-row-group scales", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#w4a16FusedVariant"], "wgslOnDisk": true } }, "w4a16_batched_f16a_f32out": { "wgsl": "fused_matmul_w4a16_batched_f16a_f32out.wgsl", "entryPoint": "main_batched_f16a_f32out", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "weightDtype": "w4a16", "variantMetadata": { "colsPerWg": 8, "tileM": 4 }, "description": "Compressed-tensors W4A16 packed batched prefill matmul with f16 activations and f32 output; keeps QAT weights packed on GPU with per-row-group scales", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#w4a16FusedVariant"], "wgslOnDisk": true } }, "q4_fused_batched": { "wgsl": "fused_matmul_q4_batched.wgsl", "entryPoint": "main_batched", "workgroup": [64,4,1], "requires": ["subgroups"], "outputDtype": "f32", "weightDtype": "q4k", "variantMetadata": { "tileM": 4 }, "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#q4kFusedVariant"], "wgslOnDisk": true } }, "q4_fused_batched_multicol_shared": { "wgsl": "fused_matmul_q4_batched_multicol_shared.wgsl", "entryPoint": "main", "workgroup": [32,4,1], "requires": [], "outputDtype": "f32", "weightDtype": "q4k", "variantMetadata": { "tileM": 4, "colsPerWg": 8 }, "description": "Shared-A Q4_K prefill matmul: 4 rows x 8 columns per workgroup while preserving block/subblock traversal", "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_multicol": { "wgsl": "fused_matmul_q4.wgsl", "entryPoint": "main_multicol", "workgroup": [256,1,1], "requires": ["subgroups"], "outputDtype": "f32", "weightDtype": "q4k", "variantMetadata": { "colsPerWg": 32 }, "description": "Multi-column GEMV for large vocab (LM head) - 32 columns per workgroup", "reachability": { "status": "pinned", "inlineConfigs": [ "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16" ], "ruleChains": ["matmul.rules.json#q4kFusedVariant"], "wgslOnDisk": true } }, "q4_fused_gemv": { "wgsl": "fused_matmul_q4.wgsl", "entryPoint": "main_gemv", "workgroup": [256,1,1], "requires": ["subgroups"], "outputDtype": "f32", "weightDtype": "q4k", "variantMetadata": { "colsPerWg": 32 }, "description": "Optimised fused Q4K GEMV - shared-A cooperative load + fast nibble extraction for maximum M=1 decode throughput without f16 precision loss", "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-embedding-0-6b-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_multicol_f16": { "wgsl": "fused_matmul_q4_multicol_f16.wgsl", "entryPoint": "main_multicol_f16", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "q4k", "variantMetadata": { "colsPerWg": 32 }, "description": "Multi-column GEMV for large vocab (LM head) - F16 output", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#q4kFusedVariant"], "wgslOnDisk": true } }, "q4_fused_multicol_f16a": { "wgsl": "fused_matmul_q4_multicol_f16a.wgsl", "entryPoint": "main_multicol_f16a", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "q4k", "variantMetadata": { "colsPerWg": 32 }, "description": "Multi-column GEMV for large vocab (LM head) - F16 input/output", "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16"], "ruleChains": ["matmul.rules.json#q4kFusedVariant"], "wgslOnDisk": true } }, "q4_fused_multicol_f16a_f32acc": { "wgsl": "fused_matmul_q4_multicol_f16a_f32acc.wgsl", "entryPoint": "main_multicol_f16a_f32acc", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "q4k", "variantMetadata": { "colsPerWg": 32 }, "description": "F16 IO + F32 accumulator twin of q4_fused_multicol_f16a. Closes the f16-accum decode-gemv gap on AMD RDNA3 where f32 accumulation outperforms f16 by ~60%.", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_metal_simd16": { "wgsl": "fused_matmul_q4_metal_simd16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["subgroups"], "outputDtype": "f32", "weightDtype": "q4k", "variantMetadata": { "colsPerWg": 16 }, "description": "Metal Q4_K GEMV with one 16-lane subgroup partition per output row.", "reachability": { "status": "pinned", "inlineConfigs": ["runtime/profiles/qwen-3-5-2b-metal-throughput"], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_batched_f16": { "wgsl": "fused_matmul_q4_batched_f16.wgsl", "entryPoint": "main_batched_f16", "workgroup": [64,4,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "q4k", "variantMetadata": { "tileM": 4 }, "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#q4kFusedVariant"], "wgslOnDisk": true } }, "q4_fused_batched_f16a": { "wgsl": "fused_matmul_q4_batched_f16a.wgsl", "entryPoint": "main_batched_f16a", "workgroup": [64,4,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "q4k", "variantMetadata": { "tileM": 4 }, "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["matmul.rules.json#q4kFusedVariant"], "wgslOnDisk": true } }, "q4_fused_batched_f16acc_f16a": { "wgsl": "fused_matmul_q4_batched_f16acc_f16a.wgsl", "entryPoint": "main_batched_f16acc_f16a", "workgroup": [256,4,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "weightDtype": "q4k", "variantMetadata": { "tileM": 4, "accumDtype": "f16" }, "description": "Batched Q4_K prefill matmul with f16 activation input, f16 accumulation, and f16 output for experimental all-f16 execution lanes.", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_prefill_tiled_f16": { "wgsl": "fused_matmul_q4_prefill_tiled_f16.wgsl", "entryPoint": "main", "workgroup": [16,16,1], "requires": ["shader-f16"], "outputDtype": "f16", "weightDtype": "q4k", "variantMetadata": { "tileM": 64, "colsPerWg": 64 }, "description": "Register-tiled Q4_K prefill matmul: 64x64 outputs per WG via 4x4 register tile per thread, amortizing Q4K dequant across 16x the outputs of q4_fused_batched_f16a.", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_widetile_f16": { "wgsl": "fused_matmul_q4_widetile_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "weightDtype": "q4k", "variantMetadata": { "tileM": 4, "colsPerWg": 256 }, "description": "WideTile Q4_K prefill matmul (adapted from ORT MatMulNBitsWideTile, MIT licensed). Each thread owns 1 output column × TILE_M rows; register-tiled accumulation. F32 A, Q4_K B (llama.cpp 256-block), F16 C. ~128x fewer WGs than q4_fused_batched_f16 at Gemma-4 prefill shape.", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_widetile_f16a": { "wgsl": "fused_matmul_q4_widetile_f16a.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "weightDtype": "q4k", "variantMetadata": { "tileM": 4, "colsPerWg": 256 }, "description": "WideTile Q4_K prefill matmul with f16 activation input and f16 output for execution-v1 f16 activation lanes.", "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16" ], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_widetile": { "wgsl": "fused_matmul_q4_widetile.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "weightDtype": "q4k", "variantMetadata": { "tileM": 4, "colsPerWg": 256 }, "description": "F32-output twin of q4_fused_widetile_f16. Same WideTile design, f32 output for callers that need it (e.g. FFN gate/up when activationDtype=f32).", "reachability": { "status": "pinned", "inlineConfigs": [ "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-embedding-0-6b-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_rmsnorm_widetile": { "wgsl": "fused_rmsnorm_q4_widetile.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "weightDtype": "q4k", "variantMetadata": { "tileM": 4, "colsPerWg": 256 }, "description": "Fused RMSNorm + Q4_K WideTile matmul. Two-pass: pass 1 computes per-row sum-of-squares across all K; pass 2 reloads A, applies normed*weight + matmul in same tileA buffer. Saves 1 standalone rmsnorm dispatch per fusion point. Bind slot 4 = norm_weight (f32).", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "q4_fused_widetile_residual": { "wgsl": "fused_matmul_q4_widetile_residual.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "weightDtype": "q4k", "variantMetadata": { "tileM": 4, "colsPerWg": 256 }, "description": "WideTile Q4_K matmul fused with residual add in one dispatch. Used at o_proj + attn_residual and ffn_down + ffn_residual call sites to cut 1 dispatch per layer per fusion point (~0.85 ms bubble each).", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "f32": { "wgsl": "matmul_f32.wgsl", "entryPoint": "main", "workgroup": [16,16,1], "requires": [], "outputDtype": "f32", "weightDtype": "f32", "reachability": { "status": "pinned", "inlineConfigs": [ "esm/esm2-t12-35m-ur50d-f32-af32", "esmc/esmc-300m-f32-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "nucleotide-transformer/nucleotide-transformer-v2-50m-f32-af32" ], "ruleChains": ["matmul.rules.json#gemvVariant","matmul.rules.json#matmulKernel"], "wgslOnDisk": true } } } }, "fused_ffn": { "description": "Fused FFN kernels (Tier 2 P0)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "gate_weights", "type": "read-only-storage" }, { "index": 3, "name": "up_weights", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 36, "fields": [ { "name": "M", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "intermediate_size", "type": "u32", "offset": 8 }, { "name": "alpha", "type": "f32", "offset": 12 }, { "name": "activation", "type": "u32", "offset": 16 } ] }, "variants": { "default": { "wgsl": "fused_ffn.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "multi": { "wgsl": "fused_ffn.wgsl", "entryPoint": "main_multi", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "f16": { "wgsl": "fused_ffn.wgsl", "entryPoint": "main_f16", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "f16_batched": { "wgsl": "fused_ffn_f16w_f32a_batched.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["subgroups"], "description": "Packed f16 gate/up weights with f32 activations and f32 output for prefill", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "f16_native": { "wgsl": "fused_ffn_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "description": "Native f16 weight reads via array bindings (no unpack2x16float)", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "f16_native_batched": { "wgsl": "fused_ffn_f16.wgsl", "entryPoint": "main_batched", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "description": "Native f16 batched variant for prefill", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "batched": { "wgsl": "fused_ffn.wgsl", "entryPoint": "main_batched", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "q4k": { "wgsl": "fused_ffn_q4k.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "q4k_metal_simd16": { "wgsl": "fused_ffn_q4k_metal_simd16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["subgroups"], "description": "Metal Q4_K fused gate/up GEMV with one 16-lane subgroup partition per output row.", "reachability": { "status": "pinned", "inlineConfigs": ["runtime/profiles/qwen-3-5-2b-metal-throughput"], "ruleChains": [], "wgslOnDisk": true } }, "q4k_f16a": { "wgsl": "fused_ffn_q4k_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "description": "Q4K fused FFN with f16 activation reads and f32 output", "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16"], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "q4k_batched": { "wgsl": "fused_ffn_q4k.wgsl", "entryPoint": "main_batched", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "q4k_batched_f16a": { "wgsl": "fused_ffn_q4k_f16.wgsl", "entryPoint": "main_batched", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "description": "Q4K fused FFN batched path with f16 activation reads and f32 output", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["fused-ffn.rules.json#variant"], "wgslOnDisk": true } }, "gate_up_gelu_gated_tiled_f16": { "wgsl": "fused_gate_up_gelu_gated_f16.wgsl", "entryPoint": "main", "workgroup": [128,1,1], "requires": ["shader-f16"], "description": "Fused gate + up FFN projections + GeGLU activation (F16 weights, per-layer intermediate_size uniform)", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "dequant": { "description": "Dequantization kernels for various quant formats", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "quantized", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "num_blocks", "type": "u32", "offset": 0 }, { "name": "block_size", "type": "u32", "offset": 4 } ] }, "variants": { "subgroup": { "wgsl": "dequant_subgroup.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": ["subgroups"], "variantMetadata": { "dispatchMode": "per_block" }, "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["dequant.rules.json#variant"], "wgslOnDisk": true } }, "subgroup_vec4": { "wgsl": "dequant_subgroup.wgsl", "entryPoint": "main_vec4", "workgroup": [64,1,1], "requires": ["subgroups"], "variantMetadata": { "dispatchMode": "per_block" }, "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["dequant.rules.json#variant"], "wgslOnDisk": true } }, "subgroup_f16out": { "wgsl": "dequant_f16_out.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "variantMetadata": { "dispatchMode": "per_block" }, "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["dequant.rules.json#variant"], "wgslOnDisk": true } }, "subgroup_vec4_f16out": { "wgsl": "dequant_f16_out_vec4.wgsl", "entryPoint": "main_vec4", "workgroup": [64,1,1], "requires": ["shader-f16"], "variantMetadata": { "dispatchMode": "per_block" }, "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["dequant.rules.json#variant"], "wgslOnDisk": true } }, "shared": { "wgsl": "dequant_shared.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "variantMetadata": { "dispatchMode": "per_block" }, "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["dequant.rules.json#variant"], "wgslOnDisk": true } }, "mxfp4": { "wgsl": "dequant_mxfp4.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "description": "MXFP4 dequantization (GPT-OSS)", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "mxfp4_expert": { "wgsl": "dequant_mxfp4_expert.wgsl", "entryPoint": "main_expert", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["dequant.rules.json#mxfp4ExpertVariant"], "wgslOnDisk": true } }, "mxfp4_expert_f16": { "wgsl": "dequant_mxfp4_expert_f16.wgsl", "entryPoint": "main_expert", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["dequant.rules.json#mxfp4ExpertVariant"], "wgslOnDisk": true } }, "q6k_f16out": { "wgsl": "dequant_q6k.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "description": "Q6_K dequantization (GGUF 6-bit quantization)", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "f32_rowwise": { "wgsl": "dequant_f32_rowwise.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "description": "Q4_K dequantization with row-wise layout (f32 output)", "uniformsOverride": { "size": 16, "fields": [ { "name": "num_blocks", "type": "u32", "offset": 0 }, { "name": "blocks_per_row", "type": "u32", "offset": 4 }, { "name": "K", "type": "u32", "offset": 8 }, { "name": "rows", "type": "u32", "offset": 12 } ] }, "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["dequant.rules.json#rowwiseVariant"], "wgslOnDisk": true } }, "f16_rowwise": { "wgsl": "dequant_f16_rowwise.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "description": "Q4_K dequantization with row-wise matrix layout for non-256-aligned K", "uniformsOverride": { "size": 16, "fields": [ { "name": "num_blocks", "type": "u32", "offset": 0 }, { "name": "blocks_per_row", "type": "u32", "offset": 4 }, { "name": "K", "type": "u32", "offset": 8 }, { "name": "rows", "type": "u32", "offset": 12 } ] }, "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["dequant.rules.json#rowwiseVariant"], "wgslOnDisk": true } }, "mxfp4_expert_fused_f16": { "wgsl": "gptoss_mxfp4_expert_fused.wgsl", "entryPoint": "main_expert", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "variantMetadata": { "modelType": "gpt-oss", "groupSize": 32, "tileShape": "vec4", "expertFormat": "mxfp4" }, "description": "GPT-OSS fused MXFP4 expert dequant kernel with vec4-oriented write path.", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["dequant.rules.json#mxfp4ExpertVariant"], "wgslOnDisk": true } } } }, "attention": { "description": "Flash attention kernels for prefill and decode", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "Q", "type": "read-only-storage" }, { "index": 2, "name": "K", "type": "read-only-storage" }, { "index": 3, "name": "V", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" }, { "index": 5, "name": "kv_len", "type": "read-only-storage" }, { "index": 6, "name": "page_table", "type": "read-only-storage" } ], "baseUniforms": { "size": 64, "fields": [ { "name": "num_heads", "type": "u32", "offset": 0 }, { "name": "num_kv_heads", "type": "u32", "offset": 4 }, { "name": "head_dim", "type": "u32", "offset": 8 }, { "name": "kv_len", "type": "u32", "offset": 12 }, { "name": "seq_len", "type": "u32", "offset": 16 }, { "name": "scale", "type": "f32", "offset": 20 }, { "name": "is_causal", "type": "u32", "offset": 24 }, { "name": "start_pos", "type": "u32", "offset": 28 }, { "name": "attn_softcap", "type": "f32", "offset": 32 }, { "name": "sliding_window", "type": "u32", "offset": 36 }, { "name": "kv_len_source", "type": "u32", "offset": 40 }, { "name": "kv_start", "type": "u32", "offset": 44 }, { "name": "page_size", "type": "u32", "offset": 48 }, { "name": "kv_layout", "type": "u32", "offset": 52 }, { "name": "_pad", "type": "u32", "offset": 56 } ] }, "variants": { "prefill": { "wgsl": "attention.wgsl", "entryPoint": "main", "workgroup": [64,1,1], "requires": [], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_large" }, "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["attention.rules.json#phase"], "wgslOnDisk": true } }, "decode": { "wgsl": "attention_decode.wgsl", "entryPoint": "attention_decode", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_large" }, "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["attention.rules.json#phase"], "wgslOnDisk": true } }, "prefill_small": { "wgsl": "attention_small.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": [], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_small" }, "reachability": { "status": "pinned", "inlineConfigs": [ "embeddinggemma/google-embeddinggemma-300m-q4k-ehf16-af32", "esm/esm2-t12-35m-ur50d-f32-af32", "esmc/esmc-300m-f32-af32", "nucleotide-transformer/nucleotide-transformer-v2-50m-f32-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "decode_small": { "wgsl": "attention_small.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": [], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_small" }, "reachability": { "status": "pinned", "inlineConfigs": [ "embeddinggemma/google-embeddinggemma-300m-q4k-ehf16-af32", "esm/esm2-t12-35m-ur50d-f32-af32", "esmc/esmc-300m-f32-af32", "nucleotide-transformer/nucleotide-transformer-v2-50m-f32-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "prefill_streaming": { "wgsl": "attention_streaming.wgsl", "entryPoint": "main", "workgroup": [1,1,1], "requires": [], "outputDtype": "f32", "variantMetadata": { "tier": "streaming" }, "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["attention.rules.json#variant"], "wgslOnDisk": true } }, "decode_streaming": { "wgsl": "attention_streaming.wgsl", "entryPoint": "main", "workgroup": [1,1,1], "requires": [], "outputDtype": "f32", "variantMetadata": { "tier": "streaming" }, "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "prefill_f16": { "wgsl": "attention_f16.wgsl", "entryPoint": "main", "workgroup": [64,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "tier": "tiled_large" }, "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "decode_f16": { "wgsl": "attention_decode_f16.wgsl", "entryPoint": "attention_decode", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "tier": "tiled_large" }, "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "prefill_small_f16": { "wgsl": "attention_small_f16.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "tier": "tiled_small" }, "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16"], "ruleChains": [], "wgslOnDisk": true } }, "decode_small_f16": { "wgsl": "attention_small_f16.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "tier": "tiled_small" }, "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16"], "ruleChains": [], "wgslOnDisk": true } }, "prefill_streaming_f16": { "wgsl": "attention_streaming_f16.wgsl", "entryPoint": "main", "workgroup": [1,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "tier": "streaming" }, "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16"], "ruleChains": [], "wgslOnDisk": true } }, "decode_streaming_f16": { "wgsl": "attention_streaming_f16.wgsl", "entryPoint": "main", "workgroup": [1,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "tier": "streaming" }, "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16"], "ruleChains": [], "wgslOnDisk": true } }, "prefill_f16kv": { "wgsl": "attention_f16kv.wgsl", "entryPoint": "main", "workgroup": [64,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_large" }, "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "decode_f16kv": { "wgsl": "attention_decode_f16kv.wgsl", "entryPoint": "attention_decode", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_large" }, "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "prefill_small_f16kv": { "wgsl": "attention_small_f16kv.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_small" }, "reachability": { "status": "pinned", "inlineConfigs": [ "amplify/amplify-120m-f16-af32", "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "decode_small_f16kv": { "wgsl": "attention_small_f16kv.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_small" }, "reachability": { "status": "pinned", "inlineConfigs": [ "amplify/amplify-120m-f16-af32", "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "prefill_head128_f16kv": { "wgsl": "attention_head128_f16kv.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_small", "queryBlockSize": 32, "exactHeadDim": 128, "maxHeadDim": 128, "requiredShared": 8192 }, "reachability": { "status": "pinned", "inlineConfigs": [ "runtime/profiles/qwen-3-reranker-0-6b-metal-stability", "runtime/profiles/qwen-3-reranker-0-6b-metal-throughput" ], "ruleChains": [], "wgslOnDisk": true } }, "prefill_head256_f16kv": { "wgsl": "attention_head256_f16kv.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_small", "queryBlockSize": 32, "exactHeadDim": 256, "maxHeadDim": 256, "requiredShared": 16384 }, "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16" ], "ruleChains": [], "wgslOnDisk": true } }, "prefill_head512_f16kv": { "wgsl": "attention_head512_f16kv.wgsl", "entryPoint": "main", "workgroup": [16,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_small", "queryBlockSize": 16, "exactHeadDim": 512, "maxHeadDim": 512, "requiredShared": 16384 }, "description": "head_dim=512 variant of prefill_head256_f16kv for Gemma-4 global attention layers. BLOCK_SIZE=16 keeps shared_block at 16KB and avoids VGPR spill from 128 vec4 private arrays per thread (iter 26 measured BLOCK_SIZE=32 at -1.7% due to register pressure).", "reachability": { "status": "pinned", "inlineConfigs": [ "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple" ], "ruleChains": [], "wgslOnDisk": true } }, "prefill_head512_f16": { "wgsl": "attention_head512_f16.wgsl", "entryPoint": "main", "workgroup": [16,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "tier": "tiled_small", "queryBlockSize": 16, "exactHeadDim": 512, "maxHeadDim": 512, "requiredShared": 16384 }, "description": "head_dim=512 pure-f16 activation prefill kernel for Gemma-4 global attention layers.", "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16" ], "ruleChains": [], "wgslOnDisk": true } }, "prefill_flash_head256_f16kv": { "wgsl": "attention_prefill_flash_head256_f16kv.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "flash_prefill", "exactHeadDim": 256, "maxHeadDim": 256, "requiredShared": 16384 }, "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "prefill_flash_reduce": { "wgsl": "attention_prefill_flash_reduce.wgsl", "entryPoint": "main", "workgroup": [64,1,1], "requires": [], "outputDtype": "f32", "variantMetadata": { "tier": "flash_reduce" }, "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "prefill_streaming_f16kv": { "wgsl": "attention_streaming_f16kv.wgsl", "entryPoint": "main", "workgroup": [1,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "streaming" }, "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "minicpm/minicpm4-0-5b-f16-af32", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-embedding-0-6b-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "decode_streaming_f16kv": { "wgsl": "attention_streaming_f16kv.wgsl", "entryPoint": "main", "workgroup": [1,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "streaming" }, "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "minicpm/minicpm4-0-5b-f16-af32", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-embedding-0-6b-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "decode_online_f16": { "wgsl": "attention_decode_online_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "variantMetadata": { "tier": "subgroup", "maxHeadDim": 512 }, "description": "Online decode kernel - f16 QKV with subgroup reductions", "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16"], "ruleChains": ["attention.rules.json#optimizedVariant"], "wgslOnDisk": true } }, "decode_online_f16kv": { "wgsl": "attention_decode_online_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f32", "variantMetadata": { "tier": "subgroup", "maxHeadDim": 512 }, "description": "Online decode kernel - f32 Q with f16 KV and subgroup reductions", "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32", "minicpm/minicpm4-0-5b-f16-af32", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-embedding-0-6b-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": ["attention.rules.json#optimizedVariant"], "wgslOnDisk": true } }, "decode_online_head256_f16kv": { "wgsl": "attention_decode_online_head256_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f32", "variantMetadata": { "tier": "subgroup", "exactHeadDim": 256, "maxHeadDim": 256 }, "description": "Online decode kernel specialized for head_dim=256 with f32 Q, f16 KV, and subgroup reductions", "reachability": { "status": "pinned", "inlineConfigs": ["gemma3/gemma-3-270m-it-f16-af32","gemma3/gemma-3-270m-it-q4k-ehf16-af32"], "ruleChains": [], "wgslOnDisk": true } }, "decode_online_head256_f16kv_output_gate": { "wgsl": "attention_decode_online_head256_f16kv_output_gate.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f32", "variantMetadata": { "tier": "subgroup", "exactHeadDim": 256, "maxHeadDim": 256 }, "description": "Online decode kernel specialized for head_dim=256 with f32 Q, f16 KV, subgroup reductions, and fused f32 output-gate sigmoid", "reachability": { "status": "pinned", "inlineConfigs": [ "runtime/profiles/qwen-3-5-0-8b-metal-parity", "runtime/profiles/qwen-3-5-0-8b-metal-throughput", "runtime/profiles/qwen-3-5-2b-metal-throughput" ], "ruleChains": [], "wgslOnDisk": true } }, "decode_chunked_f16": { "wgsl": "attention_decode_chunked_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "tier": "streaming", "maxKVLen": 2048 }, "description": "Chunked decode kernel - f16 QKV + f16 output", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["attention.rules.json#chunkedVariant"], "wgslOnDisk": true } }, "decode_chunked_f16kv": { "wgsl": "attention_decode_chunked_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "streaming", "maxKVLen": 2048 }, "description": "Chunked decode kernel - parallelizes headDim for models with few heads", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["attention.rules.json#chunkedVariant"], "wgslOnDisk": true } }, "decode_paged_f16": { "wgsl": "attention_decode_paged_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "tier": "tiled_large" }, "description": "Paged decode kernel - f16 Q/K/V with page table indirection", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["attention.rules.json#pagedVariant"], "wgslOnDisk": true } }, "decode_paged_f16kv": { "wgsl": "attention_decode_paged_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "tier": "tiled_large" }, "description": "Paged decode kernel - f16 KV with page table indirection", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["attention.rules.json#pagedVariant"], "wgslOnDisk": true } }, "decode_subgroup": { "wgsl": "attention_decode_subgroup.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["subgroups"], "outputDtype": "f32", "variantMetadata": { "tier": "subgroup" }, "description": "Subgroup-optimized decode kernel - 4 barriers (vs 80), 100% thread utilization", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["attention.rules.json#variant"], "wgslOnDisk": true } } } }, "attention_tiered": { "description": "Tiered attention kernels (hot ring + cold paged)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "Q", "type": "read-only-storage" }, { "index": 2, "name": "K_hot", "type": "read-only-storage" }, { "index": 3, "name": "V_hot", "type": "read-only-storage" }, { "index": 4, "name": "K_cold", "type": "read-only-storage" }, { "index": 5, "name": "V_cold", "type": "read-only-storage" }, { "index": 6, "name": "output", "type": "storage" }, { "index": 7, "name": "cold_page_table", "type": "read-only-storage" } ], "baseUniforms": { "size": 80, "fields": [ { "name": "num_heads", "type": "u32", "offset": 0 }, { "name": "num_kv_heads", "type": "u32", "offset": 4 }, { "name": "head_dim", "type": "u32", "offset": 8 }, { "name": "cold_len", "type": "u32", "offset": 12 }, { "name": "hot_len", "type": "u32", "offset": 16 }, { "name": "seq_len", "type": "u32", "offset": 20 }, { "name": "scale", "type": "f32", "offset": 24 }, { "name": "is_causal", "type": "u32", "offset": 28 }, { "name": "start_pos", "type": "u32", "offset": 32 }, { "name": "attn_softcap", "type": "f32", "offset": 36 }, { "name": "sliding_window", "type": "u32", "offset": 40 }, { "name": "hot_window", "type": "u32", "offset": 44 }, { "name": "hot_start", "type": "u32", "offset": 48 }, { "name": "cold_page_size", "type": "u32", "offset": 52 }, { "name": "cold_layout", "type": "u32", "offset": 56 }, { "name": "hot_layout", "type": "u32", "offset": 60 }, { "name": "_pad", "type": "u32", "offset": 64 } ] }, "variants": { "decode_tiered_f16": { "wgsl": "attention_decode_tiered_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "maxKVLen": 2048 }, "description": "Tiered decode kernel - f16 QKV + f16 output", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "decode_tiered_f16kv": { "wgsl": "attention_decode_tiered_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "maxKVLen": 2048 }, "description": "Tiered decode kernel - f32 Q + f16 KV", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "attention_tiered_quant": { "description": "Tiered attention kernels with quantized cold cache", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "Q", "type": "read-only-storage" }, { "index": 2, "name": "K_hot", "type": "read-only-storage" }, { "index": 3, "name": "V_hot", "type": "read-only-storage" }, { "index": 4, "name": "K_cold", "type": "read-only-storage" }, { "index": 5, "name": "V_cold", "type": "read-only-storage" }, { "index": 6, "name": "scales_k", "type": "read-only-storage" }, { "index": 7, "name": "scales_v", "type": "read-only-storage" }, { "index": 8, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 64, "fields": [ { "name": "num_heads", "type": "u32", "offset": 0 }, { "name": "num_kv_heads", "type": "u32", "offset": 4 }, { "name": "head_dim", "type": "u32", "offset": 8 }, { "name": "cold_len", "type": "u32", "offset": 12 }, { "name": "hot_len", "type": "u32", "offset": 16 }, { "name": "seq_len", "type": "u32", "offset": 20 }, { "name": "scale", "type": "f32", "offset": 24 }, { "name": "is_causal", "type": "u32", "offset": 28 }, { "name": "start_pos", "type": "u32", "offset": 32 }, { "name": "attn_softcap", "type": "f32", "offset": 36 }, { "name": "sliding_window", "type": "u32", "offset": 40 }, { "name": "hot_window", "type": "u32", "offset": 44 }, { "name": "hot_start", "type": "u32", "offset": 48 }, { "name": "packed_stride", "type": "u32", "offset": 52 }, { "name": "_pad", "type": "u32", "offset": 56 } ] }, "variants": { "decode_tiered_int8_f16kv": { "wgsl": "attention_decode_tiered_int8_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "maxKVLen": 2048 }, "description": "Tiered decode kernel - int8 cold K/V with f16 hot K/V", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "decode_tiered_int4_f16kv": { "wgsl": "attention_decode_tiered_int4_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "maxKVLen": 2048 }, "description": "Tiered decode kernel - int4 cold K/V with f16 hot K/V", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "decode_tiered_turboquant_f16kv": { "wgsl": "attention_decode_tiered_turboquant_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "maxKVLen": 2048 }, "description": "Tiered decode kernel - TurboQuant MSE cold K/V with f16 hot K/V", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "decode_tiered_turboquant_prod_f16kv": { "wgsl": "attention_decode_tiered_turboquant_prod_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "maxKVLen": 2048 }, "description": "Tiered decode kernel - TurboQuant prod (MSE+QJL) cold K/V with f16 hot K/V", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "attention_contiguous_quant": { "description": "Contiguous quantized attention kernels for full-attention models", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "Q", "type": "read-only-storage" }, { "index": 2, "name": "K_packed", "type": "read-only-storage" }, { "index": 3, "name": "V_packed", "type": "read-only-storage" }, { "index": 4, "name": "scales_k", "type": "read-only-storage" }, { "index": 5, "name": "scales_v", "type": "read-only-storage" }, { "index": 6, "name": "output", "type": "storage" }, { "index": 7, "name": "rotation_matrix", "type": "read-only-storage" }, { "index": 8, "name": "codebook_centroids", "type": "read-only-storage" } ], "baseUniforms": { "size": 48, "fields": [ { "name": "num_heads", "type": "u32", "offset": 0 }, { "name": "num_kv_heads", "type": "u32", "offset": 4 }, { "name": "head_dim", "type": "u32", "offset": 8 }, { "name": "kv_len", "type": "u32", "offset": 12 }, { "name": "seq_len", "type": "u32", "offset": 16 }, { "name": "scale", "type": "f32", "offset": 20 }, { "name": "is_causal", "type": "u32", "offset": 24 }, { "name": "start_pos", "type": "u32", "offset": 28 }, { "name": "attn_softcap", "type": "f32", "offset": 32 }, { "name": "sliding_window", "type": "u32", "offset": 36 }, { "name": "packed_stride", "type": "u32", "offset": 40 }, { "name": "_pad", "type": "u32", "offset": 44 } ] }, "variants": { "decode_contiguous_turboquant_f16kv": { "wgsl": "attention_decode_contiguous_turboquant_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "maxKVLen": 2048 }, "description": "Contiguous decode kernel - TurboQuant MSE quantized K/V", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "decode_contiguous_turboquant_prod_f16kv": { "wgsl": "attention_decode_contiguous_turboquant_prod_f16kv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "variantMetadata": { "maxKVLen": 2048 }, "description": "Contiguous decode kernel - TurboQuant prod (MSE+QJL) quantized K/V", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "kv_quantize": { "description": "Quantize KV cache for cold tier", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input_k", "type": "read-only-storage" }, { "index": 2, "name": "input_v", "type": "read-only-storage" }, { "index": 3, "name": "output_k", "type": "storage" }, { "index": 4, "name": "output_v", "type": "storage" }, { "index": 5, "name": "scales_k", "type": "storage" }, { "index": 6, "name": "scales_v", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "num_kv_heads", "type": "u32", "offset": 0 }, { "name": "head_dim", "type": "u32", "offset": 4 }, { "name": "start_pos", "type": "u32", "offset": 8 }, { "name": "num_tokens", "type": "u32", "offset": 12 }, { "name": "packed_stride", "type": "u32", "offset": 16 }, { "name": "_pad", "type": "u32", "offset": 20 } ] }, "variants": { "int8": { "wgsl": "kv_quantize_int8.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "description": "Quantize KV to int8 (packed u32) with per-head scale", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["kv_quantize.rules.json#variant"], "wgslOnDisk": true } }, "int4": { "wgsl": "kv_quantize_int4.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "description": "Quantize KV to int4 (packed u32) with per-head scale", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["kv_quantize.rules.json#variant"], "wgslOnDisk": true } }, "turboquant": { "wgsl": "kv_quantize_turboquant.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "description": "TurboQuant MSE: rotation + codebook quantize with per-head norm", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["kv_quantize.rules.json#variant"], "wgslOnDisk": true } }, "turboquant_prod": { "wgsl": "kv_quantize_turboquant_prod.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "description": "TurboQuant prod: MSE at (b-1) bits + 1-bit QJL residual for unbiased inner products", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["kv_quantize.rules.json#variant"], "wgslOnDisk": true } } } }, "lm_head_argmax": { "description": "Fused LM-head GEMV plus greedy argmax", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "hidden", "type": "read-only-storage" }, { "index": 2, "name": "weights", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" }, { "index": 4, "name": "temp_indices", "type": "storage" }, { "index": 5, "name": "temp_logits", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "vocab_size", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "transpose_b", "type": "u32", "offset": 8 }, { "name": "workgroups_x", "type": "u32", "offset": 12 }, { "name": "pad_token_id", "type": "u32", "offset": 16 }, { "name": "logit_softcap", "type": "f32", "offset": 20 }, { "name": "output_index", "type": "u32", "offset": 24 }, { "name": "num_groups", "type": "u32", "offset": 28 } ] }, "variants": { "phase1_f16w_f32a": { "wgsl": "lm_head_argmax_f16.wgsl", "entryPoint": "phase1", "workgroup": [256,1,1], "requires": ["shader-f16"], "weightDtype": "f16", "variantMetadata": { "colsPerWg": 64, "threadsPerCol": 4 }, "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["lm-head-argmax.rules.json#phase1Variant"], "wgslOnDisk": true } }, "phase1_q4k_f32a": { "wgsl": "lm_head_argmax_q4k.wgsl", "entryPoint": "phase1", "workgroup": [256,1,1], "weightDtype": "q4k", "variantMetadata": { "colsPerWg": 64, "threadsPerCol": 4 }, "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["lm-head-argmax.rules.json#phase1Variant"], "wgslOnDisk": true } }, "phase2": { "wgsl": "lm_head_argmax_f16.wgsl", "entryPoint": "phase2", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["lm-head-argmax.rules.json#phase2Variant"], "wgslOnDisk": true } }, "phase2_q4k": { "wgsl": "lm_head_argmax_q4k.wgsl", "entryPoint": "phase2", "workgroup": [256,1,1], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["lm-head-argmax.rules.json#phase2Variant"], "wgslOnDisk": true } } } }, "lm_head_select_logits": { "description": "Selected-token LM-head f16-weight GEMV for rerank yes/no logits", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "hidden", "type": "read-only-storage" }, { "index": 2, "name": "weights", "type": "read-only-storage" }, { "index": 3, "name": "token_ids", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "hidden_size", "type": "u32", "offset": 0 }, { "name": "vocab_size", "type": "u32", "offset": 4 }, { "name": "token_count", "type": "u32", "offset": 8 }, { "name": "hidden_offset", "type": "u32", "offset": 12 }, { "name": "transpose_b", "type": "u32", "offset": 16 }, { "name": "num_blocks_per_row", "type": "u32", "offset": 20 }, { "name": "logit_softcap", "type": "f32", "offset": 20 }, { "name": "weightDecay", "type": "f32", "offset": 24 }, { "name": "_pad1", "type": "u32", "offset": 28 } ] }, "variants": { "f16w_f32a": { "wgsl": "lm_head_select_logits.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "weightDtype": "f16", "inputDtype": "f32", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "rmsnorm": { "description": "RMS normalization kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" }, { "index": 4, "name": "residual", "type": "read-only-storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "hidden_size", "type": "u32", "offset": 0 }, { "name": "num_tokens", "type": "u32", "offset": 4 }, { "name": "eps", "type": "f32", "offset": 8 }, { "name": "has_residual", "type": "u32", "offset": 12 }, { "name": "token_stride", "type": "u32", "offset": 16 }, { "name": "output_scale", "type": "f32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "rmsnorm.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "pinned", "inlineConfigs": [ "amplify/amplify-120m-f16-af32", "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "embeddinggemma/google-embeddinggemma-300m-q4k-ehf16-af32", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32", "minicpm/minicpm4-0-5b-f16-af32", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-embedding-0-6b-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": ["rmsnorm.rules.json#variant"], "wgslOnDisk": true } }, "small": { "wgsl": "rmsnorm.wgsl", "entryPoint": "main_small", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["rmsnorm.rules.json#variant"], "wgslOnDisk": true } }, "residual": { "wgsl": "rmsnorm.wgsl", "entryPoint": "main_cached", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["rmsnorm.rules.json#variant"], "wgslOnDisk": true } }, "subgroup": { "wgsl": "rmsnorm.wgsl", "entryPoint": "main_subgroup", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["rmsnorm.rules.json#variant"], "wgslOnDisk": true } }, "small_subgroup": { "wgsl": "rmsnorm.wgsl", "entryPoint": "main_small_subgroup", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["rmsnorm.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "rmsnorm_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16"], "ruleChains": ["rmsnorm.rules.json#variant"], "wgslOnDisk": true } }, "small_f16": { "wgsl": "rmsnorm_f16.wgsl", "entryPoint": "rmsnorm_small_f16", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["rmsnorm.rules.json#variant"], "wgslOnDisk": true } } } }, "rmsnorm_pair": { "description": "Two-output sandwich RMSNorm decode kernel: post_attn_norm plus pre_ffn_norm", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "residual", "type": "read-only-storage" }, { "index": 3, "name": "post_weight", "type": "read-only-storage" }, { "index": 4, "name": "pre_weight", "type": "read-only-storage" }, { "index": 5, "name": "post_output", "type": "storage" }, { "index": 6, "name": "pre_output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "hidden_size", "type": "u32", "offset": 0 }, { "name": "num_tokens", "type": "u32", "offset": 4 }, { "name": "eps", "type": "f32", "offset": 8 }, { "name": "has_residual", "type": "u32", "offset": 12 }, { "name": "token_stride", "type": "u32", "offset": 16 }, { "name": "_pad0", "type": "u32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "rmsnorm_pair.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "subgroup": { "wgsl": "rmsnorm_pair.wgsl", "entryPoint": "main_subgroup", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "residual_rmsnorm_pair": { "description": "Two-output decode kernel: residual add plus next-layer input RMSNorm", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "residual", "type": "read-only-storage" }, { "index": 3, "name": "norm_weight", "type": "read-only-storage" }, { "index": 4, "name": "residual_output", "type": "storage" }, { "index": 5, "name": "norm_output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "hidden_size", "type": "u32", "offset": 0 }, { "name": "num_tokens", "type": "u32", "offset": 4 }, { "name": "eps", "type": "f32", "offset": 8 }, { "name": "output_scale", "type": "f32", "offset": 12 }, { "name": "token_stride", "type": "u32", "offset": 16 }, { "name": "_pad0", "type": "u32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "subgroup": { "wgsl": "residual_rmsnorm_pair.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "rmsnorm_qk": { "description": "Fused Q/K RMS normalization kernel", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "q_input", "type": "read-only-storage" }, { "index": 2, "name": "q_weight", "type": "read-only-storage" }, { "index": 3, "name": "q_output", "type": "storage" }, { "index": 4, "name": "k_input", "type": "read-only-storage" }, { "index": 5, "name": "k_weight", "type": "read-only-storage" }, { "index": 6, "name": "k_output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "q_rows", "type": "u32", "offset": 0 }, { "name": "k_rows", "type": "u32", "offset": 4 }, { "name": "head_dim", "type": "u32", "offset": 8 }, { "name": "row_stride", "type": "u32", "offset": 12 }, { "name": "eps", "type": "f32", "offset": 16 }, { "name": "_pad0", "type": "u32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "rmsnorm_qk.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["rmsnorm-qk.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "rmsnorm_qk_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["rmsnorm-qk.rules.json#variant"], "wgslOnDisk": true } } } }, "layernorm": { "description": "Layer normalization kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "bias", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "hidden_size", "type": "u32", "offset": 0 }, { "name": "num_tokens", "type": "u32", "offset": 4 }, { "name": "eps", "type": "f32", "offset": 8 }, { "name": "_pad", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "layernorm.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "pinned", "inlineConfigs": [ "esm/esm2-t12-35m-ur50d-f32-af32", "esmc/esmc-300m-f32-af32", "nucleotide-transformer/nucleotide-transformer-v2-50m-f32-af32" ], "ruleChains": ["layernorm.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "layernorm_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["layernorm.rules.json#variant"], "wgslOnDisk": true } } } }, "fused_matmul_rmsnorm": { "description": "Fused GEMV + RMSNorm for decode (M=1)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weights", "type": "read-only-storage" }, { "index": 3, "name": "norm_weight", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "M", "type": "u32", "offset": 0 }, { "name": "N", "type": "u32", "offset": 4 }, { "name": "K", "type": "u32", "offset": 8 }, { "name": "eps", "type": "f32", "offset": 12 } ] }, "variants": { "small": { "wgsl": "fused_matmul_rmsnorm.wgsl", "entryPoint": "gemv_rmsnorm_small", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["fused-matmul-rmsnorm.rules.json#variant"], "wgslOnDisk": true } }, "medium": { "wgsl": "fused_matmul_rmsnorm.wgsl", "entryPoint": "gemv_rmsnorm_medium", "workgroup": [256,1,1], "requires": [], "description": "Medium variant for N up to ~4096 (covers Gemma 3's hiddenSize=1152)", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["fused-matmul-rmsnorm.rules.json#variant"], "wgslOnDisk": true } }, "phase1": { "wgsl": "fused_matmul_rmsnorm.wgsl", "entryPoint": "gemv_rmsnorm_phase1", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "small_f16": { "wgsl": "fused_matmul_rmsnorm_f16.wgsl", "entryPoint": "gemv_rmsnorm_small", "workgroup": [256,1,1], "requires": ["shader-f16"], "description": "F16 small variant for N <= 256", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["fused-matmul-rmsnorm.rules.json#variant"], "wgslOnDisk": true } }, "medium_f16": { "wgsl": "fused_matmul_rmsnorm_f16.wgsl", "entryPoint": "gemv_rmsnorm_medium", "workgroup": [256,1,1], "requires": ["shader-f16"], "description": "F16 medium variant for N up to ~4096 (covers Gemma's hiddenSize=2304)", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["fused-matmul-rmsnorm.rules.json#variant"], "wgslOnDisk": true } } } }, "fused_matmul_residual": { "description": "Fused GEMV + Residual for decode (M=1)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weights", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" }, { "index": 4, "name": "residual", "type": "read-only-storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "M", "type": "u32", "offset": 0 }, { "name": "N", "type": "u32", "offset": 4 }, { "name": "K", "type": "u32", "offset": 8 } ] }, "variants": { "default": { "wgsl": "matmul_gemv_residual.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["fused-matmul-residual.rules.json#variant"], "wgslOnDisk": true } }, "f16": { "wgsl": "matmul_gemv_residual_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["fused-matmul-residual.rules.json#variant"], "wgslOnDisk": true } } } }, "softmax": { "description": "Softmax kernels with online normalization", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "inner_size", "type": "u32", "offset": 0 }, { "name": "outer_size", "type": "u32", "offset": 4 }, { "name": "temperature", "type": "f32", "offset": 8 }, { "name": "_pad", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "softmax.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["softmax.rules.json#variant"], "wgslOnDisk": true } }, "small": { "wgsl": "softmax.wgsl", "entryPoint": "softmax_small", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["softmax.rules.json#variant"], "wgslOnDisk": true } }, "small_subgroup": { "wgsl": "softmax_subgroup.wgsl", "entryPoint": "softmax_small_subgroup", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["softmax.rules.json#variant"], "wgslOnDisk": true } }, "subgroup": { "wgsl": "softmax_subgroup.wgsl", "entryPoint": "main_subgroup", "workgroup": [256,1,1], "requires": ["subgroups"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["softmax.rules.json#variant"], "wgslOnDisk": true } } } }, "soft_embedding": { "description": "Soft embedding helpers for block-diffusion self-conditioning", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "softmax_rows", "type": "read-only-storage" }, { "index": 2, "name": "embedding_section", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "canvas_length", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "vocab_size", "type": "u32", "offset": 8 }, { "name": "row_start", "type": "u32", "offset": 12 }, { "name": "row_count", "type": "u32", "offset": 16 }, { "name": "output_mode", "type": "u32", "offset": 20 }, { "name": "_pad0", "type": "u32", "offset": 24 }, { "name": "_pad1", "type": "u32", "offset": 28 } ] }, "variants": { "split_f16": { "wgsl": "soft_embedding_split_f16.wgsl", "entryPoint": "main", "workgroup": [8,8,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "logits_norm_stats": { "wgsl": "soft_embedding_logits_f16.wgsl", "entryPoint": "logits_norm_stats", "workgroup": [256,1,1], "requires": ["shader-f16"], "bindingsOverride": [ { "index": 1, "name": "logits", "type": "read-only-storage" }, { "index": 2, "name": "row_max", "type": "storage" }, { "index": 3, "name": "row_sum", "type": "storage" } ], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "logits_probability_chunk": { "wgsl": "soft_embedding_logits_f16.wgsl", "entryPoint": "logits_probability_chunk", "workgroup": [256,1,1], "requires": ["shader-f16"], "bindingsOverride": [ { "index": 1, "name": "logits", "type": "read-only-storage" }, { "index": 2, "name": "row_max", "type": "storage" }, { "index": 3, "name": "row_sum", "type": "storage" }, { "index": 4, "name": "probabilities", "type": "storage" } ], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "diffusion_gemma_sampling": { "description": "Block-diffusion canvas sampling stats for DiffusionGemma", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "logits", "type": "read-only-storage" }, { "index": 2, "name": "argmax_tokens", "type": "storage" }, { "index": 3, "name": "entropies", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "vocab_size", "type": "u32", "offset": 0 }, { "name": "canvas_length", "type": "u32", "offset": 4 }, { "name": "pad_token_id", "type": "u32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 }, { "name": "temperature", "type": "f32", "offset": 16 }, { "name": "logit_softcap", "type": "f32", "offset": 20 }, { "name": "_pad1", "type": "f32", "offset": 24 }, { "name": "_pad2", "type": "f32", "offset": 28 } ] }, "variants": { "entropy_stats": { "wgsl": "diffusion_gemma_sampling.wgsl", "entryPoint": "entropy_stats", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "rope": { "description": "Rotary position embeddings kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "storage" }, { "index": 2, "name": "freqs_cos", "type": "storage" }, { "index": 3, "name": "freqs_sin", "type": "storage" } ], "baseUniforms": { "size": 48, "fields": [ { "name": "seq_len", "type": "u32", "offset": 0 }, { "name": "num_heads", "type": "u32", "offset": 4 }, { "name": "head_dim", "type": "u32", "offset": 8 }, { "name": "start_pos", "type": "u32", "offset": 12 }, { "name": "rope_base", "type": "f32", "offset": 16 }, { "name": "rope_scale", "type": "f32", "offset": 20 }, { "name": "rotary_dim", "type": "u32", "offset": 24 }, { "name": "interleaved", "type": "u32", "offset": 28 }, { "name": "pair_span_dim", "type": "u32", "offset": 32 }, { "name": "_pad0", "type": "u32", "offset": 36 }, { "name": "_pad1", "type": "u32", "offset": 40 }, { "name": "_pad2", "type": "u32", "offset": 44 } ] }, "variants": { "default": { "wgsl": "rope.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "pinned", "inlineConfigs": [ "amplify/amplify-120m-f16-af32", "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "embeddinggemma/google-embeddinggemma-300m-q4k-ehf16-af32", "esm/esm2-t12-35m-ur50d-f32-af32", "esmc/esmc-300m-f32-af32", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32", "minicpm/minicpm4-0-5b-f16-af32", "nucleotide-transformer/nucleotide-transformer-v2-50m-f32-af32", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-embedding-0-6b-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": ["rope.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "rope_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16"], "ruleChains": ["rope.rules.json#variant"], "wgslOnDisk": true } }, "compute_freqs": { "wgsl": "rope.wgsl", "entryPoint": "rope_compute_freqs", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "compute_freqs_f16": { "wgsl": "rope_f16.wgsl", "entryPoint": "rope_compute_freqs", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "qk": { "wgsl": "rope.wgsl", "entryPoint": "rope_qk", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "qk_f16": { "wgsl": "rope_f16.wgsl", "entryPoint": "rope_qk", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "ntk": { "wgsl": "rope.wgsl", "entryPoint": "rope_ntk_scaled", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "ntk_f16": { "wgsl": "rope_f16.wgsl", "entryPoint": "rope_ntk_scaled", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "yarn": { "wgsl": "rope.wgsl", "entryPoint": "rope_yarn", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "yarn_f16": { "wgsl": "rope_f16.wgsl", "entryPoint": "rope_yarn", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "rope_qk": { "description": "Rotary position embeddings for separate Q/K buffers", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "q_input", "type": "storage" }, { "index": 2, "name": "k_input", "type": "storage" }, { "index": 3, "name": "freqs_cos", "type": "read-only-storage" }, { "index": 4, "name": "freqs_sin", "type": "read-only-storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "seq_len", "type": "u32", "offset": 0 }, { "name": "num_q_heads", "type": "u32", "offset": 4 }, { "name": "num_kv_heads", "type": "u32", "offset": 8 }, { "name": "head_dim", "type": "u32", "offset": 12 }, { "name": "start_pos", "type": "u32", "offset": 16 }, { "name": "rotary_dim", "type": "u32", "offset": 20 }, { "name": "interleaved", "type": "u32", "offset": 24 }, { "name": "pair_span_dim", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "rope_qk.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["rope-qk.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "rope_qk_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["rope-qk.rules.json#variant"], "wgslOnDisk": true } } } }, "silu": { "description": "SiLU activation kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" }, { "index": 3, "name": "gate", "type": "read-only-storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "rowsplit_dim", "type": "u32", "offset": 4 } ] }, "variants": { "default": { "wgsl": "silu.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "pinned", "inlineConfigs": [ "amplify/amplify-120m-f16-af32", "esmc/esmc-300m-f32-af32", "minicpm/minicpm4-0-5b-f16-af32", "nucleotide-transformer/nucleotide-transformer-v2-50m-f32-af32", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-embedding-0-6b-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": ["silu.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "silu_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["silu.rules.json#variant"], "wgslOnDisk": true } } } }, "gelu": { "description": "GeLU/GeGLU activation kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" }, { "index": 3, "name": "gate", "type": "read-only-storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "rowsplit_dim", "type": "u32", "offset": 4 } ] }, "variants": { "gelu": { "wgsl": "gelu.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "embeddinggemma/google-embeddinggemma-300m-q4k-ehf16-af32", "esm/esm2-t12-35m-ur50d-f32-af32", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32" ], "ruleChains": ["gelu.rules.json#variant"], "wgslOnDisk": true } }, "gelu_f16": { "wgsl": "gelu_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16"], "ruleChains": ["gelu.rules.json#variant"], "wgslOnDisk": true } } } }, "scale": { "description": "Tensor scaling kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "scale", "type": "f32", "offset": 4 } ] }, "variants": { "default": { "wgsl": "scale.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["scale.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "scale_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["scale.rules.json#variant"], "wgslOnDisk": true } }, "inplace": { "wgsl": "scale.wgsl", "entryPoint": "main_inplace", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["scale.rules.json#variant"], "wgslOnDisk": true } }, "inplace_f16": { "wgsl": "scale_f16.wgsl", "entryPoint": "main_inplace", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["scale.rules.json#variant"], "wgslOnDisk": true } } } }, "clamp": { "description": "Tensor clamp kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "data", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "_pad0", "type": "u32", "offset": 4 }, { "name": "min", "type": "f32", "offset": 8 }, { "name": "max", "type": "f32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "clamp.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "activation_static_qdq": { "description": "Static activation quantize/dequantize kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "_pad0", "type": "u32", "offset": 4 }, { "name": "scale", "type": "f32", "offset": 8 }, { "name": "invScale", "type": "f32", "offset": 12 }, { "name": "qmin", "type": "f32", "offset": 16 }, { "name": "qmax", "type": "f32", "offset": 20 } ] }, "variants": { "default": { "wgsl": "activation_static_qdq.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "energy_eval": { "description": "Energy evaluation kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "state", "type": "read-only-storage" }, { "index": 2, "name": "target", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "count", "type": "u32", "offset": 0 }, { "name": "scale", "type": "f32", "offset": 4 } ] }, "variants": { "eval": { "wgsl": "energy_eval.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["energy.rules.json#evalVariant"], "wgslOnDisk": true } }, "eval_f16": { "wgsl": "energy_eval_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["energy.rules.json#evalVariant"], "wgslOnDisk": true } } } }, "energy_update": { "description": "Energy update kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "state", "type": "storage" }, { "index": 2, "name": "target", "type": "read-only-storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "count", "type": "u32", "offset": 0 }, { "name": "stepSize", "type": "f32", "offset": 4 }, { "name": "gradientScale", "type": "f32", "offset": 8 } ] }, "variants": { "update": { "wgsl": "energy_update.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["energy.rules.json#updateVariant"], "wgslOnDisk": true } }, "update_f16": { "wgsl": "energy_update_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["energy.rules.json#updateVariant"], "wgslOnDisk": true } } } }, "energy_quintel_update": { "description": "Quintel energy update kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "state", "type": "storage" } ], "baseUniforms": { "size": 64, "fields": [ { "name": "count", "type": "u32", "offset": 0 }, { "name": "size", "type": "u32", "offset": 4 }, { "name": "flags", "type": "u32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 }, { "name": "stepSize", "type": "f32", "offset": 16 }, { "name": "gradientScale", "type": "f32", "offset": 20 }, { "name": "countDiff", "type": "f32", "offset": 24 }, { "name": "centerTarget", "type": "f32", "offset": 28 }, { "name": "symmetryWeight", "type": "f32", "offset": 32 }, { "name": "countWeight", "type": "f32", "offset": 36 }, { "name": "centerWeight", "type": "f32", "offset": 40 }, { "name": "binarizeWeight", "type": "f32", "offset": 44 }, { "name": "clampMin", "type": "f32", "offset": 48 }, { "name": "clampMax", "type": "f32", "offset": 52 }, { "name": "_pad1", "type": "f32", "offset": 56 }, { "name": "_pad2", "type": "f32", "offset": 60 } ] }, "variants": { "quintel_update": { "wgsl": "energy_quintel_update.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["energy.rules.json#quintelUpdateVariant"], "wgslOnDisk": true } }, "quintel_update_f16": { "wgsl": "energy_quintel_update_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["energy.rules.json#quintelUpdateVariant"], "wgslOnDisk": true } } } }, "energy_quintel_grad": { "description": "Quintel energy gradient kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "state", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 64, "fields": [ { "name": "count", "type": "u32", "offset": 0 }, { "name": "size", "type": "u32", "offset": 4 }, { "name": "flags", "type": "u32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 }, { "name": "stepSize", "type": "f32", "offset": 16 }, { "name": "gradientScale", "type": "f32", "offset": 20 }, { "name": "countDiff", "type": "f32", "offset": 24 }, { "name": "centerTarget", "type": "f32", "offset": 28 }, { "name": "symmetryWeight", "type": "f32", "offset": 32 }, { "name": "countWeight", "type": "f32", "offset": 36 }, { "name": "centerWeight", "type": "f32", "offset": 40 }, { "name": "binarizeWeight", "type": "f32", "offset": 44 }, { "name": "clampMin", "type": "f32", "offset": 48 }, { "name": "clampMax", "type": "f32", "offset": 52 }, { "name": "_pad1", "type": "f32", "offset": 56 }, { "name": "_pad2", "type": "f32", "offset": 60 } ] }, "variants": { "quintel_grad": { "wgsl": "energy_quintel_grad.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["energy.rules.json#quintelGradVariant"], "wgslOnDisk": true } }, "quintel_grad_f16": { "wgsl": "energy_quintel_grad_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["energy.rules.json#quintelGradVariant"], "wgslOnDisk": true } } } }, "energy_quintel_reduce": { "description": "Quintel energy reduction kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "state", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 48, "fields": [ { "name": "count", "type": "u32", "offset": 0 }, { "name": "size", "type": "u32", "offset": 4 }, { "name": "flags", "type": "u32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 }, { "name": "symmetryWeight", "type": "f32", "offset": 16 }, { "name": "centerWeight", "type": "f32", "offset": 20 }, { "name": "binarizeWeight", "type": "f32", "offset": 24 }, { "name": "centerTarget", "type": "f32", "offset": 28 }, { "name": "_pad1", "type": "f32", "offset": 32 }, { "name": "_pad2", "type": "f32", "offset": 36 }, { "name": "_pad3", "type": "f32", "offset": 40 }, { "name": "_pad4", "type": "f32", "offset": 44 } ] }, "variants": { "quintel_reduce": { "wgsl": "energy_quintel_reduce.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["energy.rules.json#quintelReduceVariant"], "wgslOnDisk": true } }, "quintel_reduce_f16": { "wgsl": "energy_quintel_reduce_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["energy.rules.json#quintelReduceVariant"], "wgslOnDisk": true } } } }, "gather": { "description": "Embedding lookup kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "indices", "type": "read-only-storage" }, { "index": 2, "name": "embeddings", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 36, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "vocab_size", "type": "u32", "offset": 8 }, { "name": "transpose", "type": "u32", "offset": 12 }, { "name": "index_offset", "type": "u32", "offset": 16 }, { "name": "input_hidden_size", "type": "u32", "offset": 20 }, { "name": "hidden_offset", "type": "u32", "offset": 24 }, { "name": "_pad0", "type": "u32", "offset": 28 }, { "name": "_pad1", "type": "u32", "offset": 32 } ] }, "variants": { "default": { "wgsl": "gather.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "embeddinggemma/google-embeddinggemma-300m-q4k-ehf16-af32", "esm/esm2-t12-35m-ur50d-f32-af32", "esmc/esmc-300m-f32-af32", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "nucleotide-transformer/nucleotide-transformer-v2-50m-f32-af32" ], "ruleChains": ["gather.rules.json#variant"], "wgslOnDisk": true } }, "vec4": { "wgsl": "gather_vec4.wgsl", "entryPoint": "gather_vec4", "workgroup": [64,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["gather.rules.json#variant"], "wgslOnDisk": true } }, "f16": { "wgsl": "gather_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "description": "F16 embeddings -> F32 output (for weight-tied lm_head optimization)", "reachability": { "status": "pinned", "inlineConfigs": [ "amplify/amplify-120m-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32", "minicpm/minicpm4-0-5b-f16-af32", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-embedding-0-6b-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": ["gather.rules.json#variant"], "wgslOnDisk": true } }, "f16_vec4": { "wgsl": "gather_f16_vec4.wgsl", "entryPoint": "gather_vec4", "workgroup": [64,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["gather.rules.json#variant"], "wgslOnDisk": true } }, "f16_out": { "wgsl": "gather_f16_out.wgsl", "entryPoint": "gather_f16_out", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "outputBinding": 4 }, "description": "F32 embeddings -> F16 output", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["gather.rules.json#variant"], "wgslOnDisk": true } }, "vec4_f16_out": { "wgsl": "gather_vec4_f16_out.wgsl", "entryPoint": "gather_vec4_f16_out", "workgroup": [64,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "outputBinding": 4 }, "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["gather.rules.json#variant"], "wgslOnDisk": true } }, "f16_f16_out": { "wgsl": "gather_f16_f16_out.wgsl", "entryPoint": "gather_f16_out", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "outputBinding": 4 }, "description": "F16 embeddings -> F16 output", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["gather.rules.json#variant"], "wgslOnDisk": true } }, "f16_vec4_f16_out": { "wgsl": "gather_f16_vec4_f16_out.wgsl", "entryPoint": "gather_vec4_f16_out", "workgroup": [64,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "outputBinding": 4 }, "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["gather.rules.json#variant"], "wgslOnDisk": true } }, "litert_int4_vec4_f16_out": { "wgsl": "gather_litert_int4_vec4_f16_out.wgsl", "entryPoint": "gather_litert_int4_vec4_f16_out", "workgroup": [64,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "outputBinding": 4 }, "description": "LiteRT signed INT4 affine embeddings -> F16 output", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["gather.rules.json#variant"], "wgslOnDisk": true } }, "litert_int4_vec4_f32_out": { "wgsl": "gather_litert_int4_vec4_f32_out.wgsl", "entryPoint": "gather_litert_int4_vec4_f32_out", "workgroup": [64,1,1], "requires": [], "outputDtype": "f32", "variantMetadata": { "outputBinding": 4 }, "description": "LiteRT signed INT4 affine embeddings -> F32 output", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["gather.rules.json#variant"], "wgslOnDisk": true } } } }, "gather_split4": { "description": "Row-split embedding lookup kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "indices", "type": "read-only-storage" }, { "index": 2, "name": "embeddings0", "type": "read-only-storage" }, { "index": 3, "name": "embeddings1", "type": "read-only-storage" }, { "index": 4, "name": "embeddings2", "type": "read-only-storage" }, { "index": 5, "name": "embeddings3", "type": "read-only-storage" }, { "index": 6, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 48, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "vocab_size", "type": "u32", "offset": 8 }, { "name": "index_offset", "type": "u32", "offset": 12 }, { "name": "input_hidden_size", "type": "u32", "offset": 16 }, { "name": "hidden_offset", "type": "u32", "offset": 20 }, { "name": "section0_rows", "type": "u32", "offset": 24 }, { "name": "section1_rows", "type": "u32", "offset": 28 }, { "name": "section2_rows", "type": "u32", "offset": 32 }, { "name": "section3_rows", "type": "u32", "offset": 36 }, { "name": "_pad0", "type": "u32", "offset": 40 }, { "name": "_pad1", "type": "u32", "offset": 44 } ] }, "variants": { "f16_vec4_f16_out": { "wgsl": "gather_split4_f16_vec4_f16_out.wgsl", "entryPoint": "gather_vec4_f16_out", "workgroup": [64,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "description": "F16 split embeddings -> F16 output", "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16"], "ruleChains": [], "wgslOnDisk": true } } } }, "gather_split8": { "description": "Row-split embedding lookup kernels with up to eight sections", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "indices", "type": "read-only-storage" }, { "index": 2, "name": "embeddings0", "type": "read-only-storage" }, { "index": 3, "name": "embeddings1", "type": "read-only-storage" }, { "index": 4, "name": "embeddings2", "type": "read-only-storage" }, { "index": 5, "name": "embeddings3", "type": "read-only-storage" }, { "index": 6, "name": "embeddings4", "type": "read-only-storage" }, { "index": 7, "name": "embeddings5", "type": "read-only-storage" }, { "index": 8, "name": "embeddings6", "type": "read-only-storage" }, { "index": 9, "name": "embeddings7", "type": "read-only-storage" }, { "index": 10, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 64, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "vocab_size", "type": "u32", "offset": 8 }, { "name": "index_offset", "type": "u32", "offset": 12 }, { "name": "input_hidden_size", "type": "u32", "offset": 16 }, { "name": "hidden_offset", "type": "u32", "offset": 20 }, { "name": "section0_rows", "type": "u32", "offset": 24 }, { "name": "section1_rows", "type": "u32", "offset": 28 }, { "name": "section2_rows", "type": "u32", "offset": 32 }, { "name": "section3_rows", "type": "u32", "offset": 36 }, { "name": "section4_rows", "type": "u32", "offset": 40 }, { "name": "section5_rows", "type": "u32", "offset": 44 }, { "name": "section6_rows", "type": "u32", "offset": 48 }, { "name": "section7_rows", "type": "u32", "offset": 52 }, { "name": "_pad0", "type": "u32", "offset": 56 }, { "name": "_pad1", "type": "u32", "offset": 60 } ] }, "variants": { "f16_vec4_f16_out": { "wgsl": "gather_split8_f16_vec4_f16_out.wgsl", "entryPoint": "gather_vec4_f16_out", "workgroup": [64,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "description": "F16 split8 embeddings -> F16 output", "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple"], "ruleChains": [], "wgslOnDisk": true } }, "f16_vec4_f32_out": { "wgsl": "gather_split8_f16_vec4_f32_out.wgsl", "entryPoint": "gather_vec4_f32_out", "workgroup": [64,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "description": "F16 split8 embeddings -> F32 output", "reachability": { "status": "pinned", "inlineConfigs": [ "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple" ], "ruleChains": [], "wgslOnDisk": true } } } }, "residual": { "description": "Residual addition kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "residual", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "scale", "type": "f32", "offset": 4 } ] }, "variants": { "default": { "wgsl": "residual.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "pinned", "inlineConfigs": [ "amplify/amplify-120m-f16-af32", "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "embeddinggemma/google-embeddinggemma-300m-q4k-ehf16-af32", "esm/esm2-t12-35m-ur50d-f32-af32", "esmc/esmc-300m-f32-af32", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32", "minicpm/minicpm4-0-5b-f16-af32", "nucleotide-transformer/nucleotide-transformer-v2-50m-f32-af32", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-embedding-0-6b-q4k-ehf16-af32", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": ["residual.rules.json#biasAddVariant","residual.rules.json#residualVariant"], "wgslOnDisk": true } }, "vec4": { "wgsl": "residual_vec4.wgsl", "entryPoint": "add_vec4", "workgroup": [64,1,1], "requires": [], "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["residual.rules.json#residualVariant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "residual_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "pinned", "inlineConfigs": ["gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16"], "ruleChains": ["residual.rules.json#residualVariant"], "wgslOnDisk": true } }, "vec4_f16": { "wgsl": "residual_f16_vec4.wgsl", "entryPoint": "add_vec4", "workgroup": [64,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "model-selectable", "inlineConfigs": [], "ruleChains": ["residual.rules.json#residualVariant"], "wgslOnDisk": true } } } }, "topk": { "description": "Top-K selection kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "values", "type": "storage" }, { "index": 3, "name": "indices", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "k", "type": "u32", "offset": 4 } ] }, "variants": { "default": { "wgsl": "topk.wgsl", "entryPoint": "main", "workgroup": [32,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "small": { "wgsl": "topk.wgsl", "entryPoint": "topk_2_small", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "fused": { "wgsl": "topk.wgsl", "entryPoint": "softmax_topk", "workgroup": [32,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "fused_f16": { "wgsl": "topk_f16.wgsl", "entryPoint": "softmax_topk", "workgroup": [32,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "fused_f16_w16": { "wgsl": "topk_f16_weights.wgsl", "entryPoint": "softmax_topk", "workgroup": [32,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "gptoss_router_topk": { "wgsl": "gptoss_router_topk.wgsl", "entryPoint": "softmax_topk", "workgroup": [256,1,1], "requires": ["shader-f16","subgroups"], "outputDtype": "f16", "variantMetadata": { "modelType": "gpt-oss", "routerDtype": "f16", "weightsDtype": "f16", "tileShape": "token-major", "groupSize": 32 }, "description": "GPT-OSS deterministic router top-k kernel (f16 logits/weights).", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "scatter_add": { "description": "MoE output combination kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "indices", "type": "read-only-storage" }, { "index": 3, "name": "weights", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "top_k", "type": "u32", "offset": 8 }, { "name": "num_experts", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "scatter_add.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "vec4": { "wgsl": "scatter_add_vec4.wgsl", "entryPoint": "scatter_add_vec4", "workgroup": [64,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "dynamic": { "wgsl": "scatter_add_dynamic.wgsl", "entryPoint": "scatter_add_dynamic", "workgroup": [256,1,1], "requires": [], "bindingsOverride": [ { "index": 2, "name": "routing_indices", "type": "read-only-storage" }, { "index": 3, "name": "routing_weights", "type": "read-only-storage" }, { "index": 4, "name": "token_offsets", "type": "read-only-storage" }, { "index": 5, "name": "output", "type": "storage" } ], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "dynamic_expert_scale": { "wgsl": "scatter_add_dynamic_expert_scale.wgsl", "entryPoint": "scatter_add_dynamic", "workgroup": [256,1,1], "requires": [], "bindingsOverride": [ { "index": 2, "name": "routing_indices", "type": "read-only-storage" }, { "index": 3, "name": "routing_weights", "type": "read-only-storage" }, { "index": 4, "name": "token_offsets", "type": "read-only-storage" }, { "index": 5, "name": "expert_scales", "type": "read-only-storage" }, { "index": 6, "name": "output", "type": "storage" } ], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "dynamic_f16": { "wgsl": "scatter_add_dynamic_f16.wgsl", "entryPoint": "scatter_add_dynamic", "workgroup": [256,1,1], "requires": ["shader-f16"], "bindingsOverride": [ { "index": 2, "name": "routing_indices", "type": "read-only-storage" }, { "index": 3, "name": "routing_weights", "type": "read-only-storage" }, { "index": 4, "name": "token_offsets", "type": "read-only-storage" }, { "index": 5, "name": "output", "type": "storage" } ], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "dynamic_f16_expert_scale": { "wgsl": "scatter_add_dynamic_f16_expert_scale.wgsl", "entryPoint": "scatter_add_dynamic", "workgroup": [256,1,1], "requires": ["shader-f16"], "bindingsOverride": [ { "index": 2, "name": "routing_indices", "type": "read-only-storage" }, { "index": 3, "name": "routing_weights", "type": "read-only-storage" }, { "index": 4, "name": "token_offsets", "type": "read-only-storage" }, { "index": 5, "name": "expert_scales", "type": "read-only-storage" }, { "index": 6, "name": "output", "type": "storage" } ], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "dynamic_f16_w16": { "wgsl": "scatter_add_dynamic_f16_weights.wgsl", "entryPoint": "scatter_add_dynamic", "workgroup": [256,1,1], "requires": ["shader-f16"], "bindingsOverride": [ { "index": 2, "name": "routing_indices", "type": "read-only-storage" }, { "index": 3, "name": "routing_weights", "type": "read-only-storage" }, { "index": 4, "name": "token_offsets", "type": "read-only-storage" }, { "index": 5, "name": "output", "type": "storage" } ], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "dynamic_f16_w16_expert_scale": { "wgsl": "scatter_add_dynamic_f16_weights_expert_scale.wgsl", "entryPoint": "scatter_add_dynamic", "workgroup": [256,1,1], "requires": ["shader-f16"], "bindingsOverride": [ { "index": 2, "name": "routing_indices", "type": "read-only-storage" }, { "index": 3, "name": "routing_weights", "type": "read-only-storage" }, { "index": 4, "name": "token_offsets", "type": "read-only-storage" }, { "index": 5, "name": "expert_scales", "type": "read-only-storage" }, { "index": 6, "name": "output", "type": "storage" } ], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "accumulate": { "wgsl": "scatter_add.wgsl", "entryPoint": "scatter_add_accumulate", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "gemma4_route_q4_matmul": { "description": "Route-wise Gemma4 packed Q4_K expert matmul kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "route_indices", "type": "read-only-storage" }, { "index": 3, "name": "weights", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "num_routes", "type": "u32", "offset": 0 }, { "name": "top_k", "type": "u32", "offset": 4 }, { "name": "N", "type": "u32", "offset": 8 }, { "name": "K", "type": "u32", "offset": 12 }, { "name": "num_blocks_per_row", "type": "u32", "offset": 16 }, { "name": "input_mode", "type": "u32", "offset": 20 }, { "name": "alpha", "type": "f32", "offset": 24 }, { "name": "_pad0", "type": "u32", "offset": 28 } ] }, "variants": { "f16a": { "wgsl": "gemma4_route_q4_matmul_f16a.wgsl", "entryPoint": "route_matmul_f16a", "workgroup": [64,4,1], "requires": ["shader-f16","subgroups"], "wgslOverrides": { "TILE_M": 4, "THREADS_PER_COL": 64 }, "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "scatter_add_routes": { "description": "Route-wise MoE output combination kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "route_outputs", "type": "read-only-storage" }, { "index": 2, "name": "routing_indices", "type": "read-only-storage" }, { "index": 3, "name": "routing_weights", "type": "read-only-storage" }, { "index": 4, "name": "expert_scales", "type": "read-only-storage" }, { "index": 5, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "top_k", "type": "u32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 } ] }, "variants": { "f16_w16_expert_scale": { "wgsl": "scatter_add_routes_f16_weights_expert_scale.wgsl", "entryPoint": "scatter_add_routes", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "moe_gather": { "description": "Token gathering by expert kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "expert_assignments", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "num_experts", "type": "u32", "offset": 8 }, { "name": "top_k", "type": "u32", "offset": 12 }, { "name": "max_tokens_per_expert", "type": "u32", "offset": 16 }, { "name": "threads_per_row", "type": "u32", "offset": 20 }, { "name": "_pad2", "type": "u32", "offset": 24 }, { "name": "_pad3", "type": "u32", "offset": 28 } ] }, "variants": { "count": { "wgsl": "moe_gather.wgsl", "entryPoint": "count_and_map", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "count_f16": { "wgsl": "moe_gather_f16.wgsl", "entryPoint": "count_and_map", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "gather": { "wgsl": "moe_gather.wgsl", "entryPoint": "gather_tokens", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "gather_f16": { "wgsl": "moe_gather_f16.wgsl", "entryPoint": "gather_tokens", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "single_pass": { "wgsl": "moe_gather.wgsl", "entryPoint": "gather_single_pass", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "moe_offsets": { "description": "Build routing token offsets for dynamic MoE scatter on GPU", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "token_counts", "type": "read-only-storage" }, { "index": 2, "name": "token_map", "type": "read-only-storage" }, { "index": 3, "name": "token_offsets", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "num_experts", "type": "u32", "offset": 4 }, { "name": "top_k", "type": "u32", "offset": 8 }, { "name": "max_tokens_per_expert", "type": "u32", "offset": 12 }, { "name": "_pad0", "type": "u32", "offset": 16 }, { "name": "_pad1", "type": "u32", "offset": 20 }, { "name": "_pad2", "type": "u32", "offset": 24 }, { "name": "_pad3", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "moe_offsets.wgsl", "entryPoint": "build_offsets", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "swiglu": { "description": "SwiGLU activation (gated) kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "gate", "type": "read-only-storage" }, { "index": 2, "name": "up", "type": "read-only-storage" }, { "index": 3, "name": "bias", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 } ] }, "variants": { "rowsplit_bias": { "wgsl": "fused_swiglu.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "rowsplit_bias_f16": { "wgsl": "fused_swiglu_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "bias_add": { "description": "Bias addition kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "data", "type": "storage" }, { "index": 2, "name": "bias", "type": "read-only-storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "dim", "type": "u32", "offset": 4 }, { "name": "data_offset", "type": "u32", "offset": 8 }, { "name": "bias_offset", "type": "u32", "offset": 12 }, { "name": "token_stride", "type": "u32", "offset": 16 }, { "name": "_pad0", "type": "u32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "bias_add.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "pinned", "inlineConfigs": [ "esm/esm2-t12-35m-ur50d-f32-af32", "esmc/esmc-300m-f32-af32", "nucleotide-transformer/nucleotide-transformer-v2-50m-f32-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "f16": { "wgsl": "bias_add_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "cast": { "description": "Type casting kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 } ] }, "variants": { "f32_to_f16": { "wgsl": "cast_f32_to_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "pinned", "inlineConfigs": ["diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16"], "ruleChains": [], "wgslOnDisk": true } }, "f16_to_f32": { "wgsl": "cast_f16_to_f32.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "kv_cache_write": { "description": "Direct KV cache write kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input_k", "type": "read-only-storage" }, { "index": 2, "name": "input_v", "type": "read-only-storage" }, { "index": 3, "name": "output_k", "type": "storage" }, { "index": 4, "name": "output_v", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "src_offset", "type": "u32", "offset": 0 }, { "name": "dst_offset", "type": "u32", "offset": 4 }, { "name": "count", "type": "u32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 } ] }, "variants": { "f32_to_f16": { "wgsl": "kv_cache_write_f32_to_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "split_qkv": { "description": "Split fused QKV output into separate Q, K, V buffers", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "fused_qkv", "type": "read-only-storage" }, { "index": 2, "name": "Q", "type": "storage" }, { "index": 3, "name": "K", "type": "storage" }, { "index": 4, "name": "V", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "q_size", "type": "u32", "offset": 4 }, { "name": "k_size", "type": "u32", "offset": 8 }, { "name": "v_size", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "split_qkv.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["split-qkv.rules.json#variant"], "wgslOnDisk": true } }, "f16": { "wgsl": "split_qkv_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["split-qkv.rules.json#variant"], "wgslOnDisk": true } } } }, "split_qkv_rmsnorm_qk": { "description": "Split fused QKV output and apply weighted RMSNorm to Q and K buffers", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "fused_qkv", "type": "read-only-storage" }, { "index": 2, "name": "q_weight", "type": "read-only-storage" }, { "index": 3, "name": "k_weight", "type": "read-only-storage" }, { "index": 4, "name": "Q", "type": "storage" }, { "index": 5, "name": "K", "type": "storage" }, { "index": 6, "name": "V", "type": "storage" } ], "baseUniforms": { "size": 48, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "q_size", "type": "u32", "offset": 4 }, { "name": "k_size", "type": "u32", "offset": 8 }, { "name": "v_size", "type": "u32", "offset": 12 }, { "name": "num_heads", "type": "u32", "offset": 16 }, { "name": "num_kv_heads", "type": "u32", "offset": 20 }, { "name": "head_dim", "type": "u32", "offset": 24 }, { "name": "workgroup_stride", "type": "u32", "offset": 28 }, { "name": "qk_rows", "type": "u32", "offset": 32 }, { "name": "total_v", "type": "u32", "offset": 36 }, { "name": "eps", "type": "f32", "offset": 40 }, { "name": "_pad0", "type": "u32", "offset": 44 } ] }, "variants": { "default": { "wgsl": "split_qkv_rmsnorm_qk.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "split_qkv_rmsnorm_rope_qk": { "description": "Split fused QKV output and apply weighted RMSNorm plus partial/full RoPE to Q and K buffers", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "fused_qkv", "type": "read-only-storage" }, { "index": 2, "name": "q_weight", "type": "read-only-storage" }, { "index": 3, "name": "k_weight", "type": "read-only-storage" }, { "index": 4, "name": "freqs_cos", "type": "read-only-storage" }, { "index": 5, "name": "freqs_sin", "type": "read-only-storage" }, { "index": 6, "name": "Q", "type": "storage" }, { "index": 7, "name": "K", "type": "storage" }, { "index": 8, "name": "V", "type": "storage" } ], "baseUniforms": { "size": 64, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "q_size", "type": "u32", "offset": 4 }, { "name": "k_size", "type": "u32", "offset": 8 }, { "name": "v_size", "type": "u32", "offset": 12 }, { "name": "num_heads", "type": "u32", "offset": 16 }, { "name": "num_kv_heads", "type": "u32", "offset": 20 }, { "name": "head_dim", "type": "u32", "offset": 24 }, { "name": "workgroup_stride", "type": "u32", "offset": 28 }, { "name": "qk_rows", "type": "u32", "offset": 32 }, { "name": "total_v", "type": "u32", "offset": 36 }, { "name": "start_pos", "type": "u32", "offset": 40 }, { "name": "rotary_dim", "type": "u32", "offset": 44 }, { "name": "eps", "type": "f32", "offset": 48 }, { "name": "interleaved", "type": "u32", "offset": 52 }, { "name": "_pad1", "type": "u32", "offset": 56 }, { "name": "_pad2", "type": "u32", "offset": 60 } ] }, "variants": { "default": { "wgsl": "split_qkv_rmsnorm_rope_qk.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "split_qkv_rmsnorm_rope_qk_f16kv_cache": { "description": "Split fused QKV output, apply weighted RMSNorm plus partial/full RoPE to Q/K, and write K/V directly to contiguous f16 KV cache", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "fused_qkv", "type": "read-only-storage" }, { "index": 2, "name": "q_weight", "type": "read-only-storage" }, { "index": 3, "name": "k_weight", "type": "read-only-storage" }, { "index": 4, "name": "freqs_cos", "type": "read-only-storage" }, { "index": 5, "name": "freqs_sin", "type": "read-only-storage" }, { "index": 6, "name": "Q", "type": "storage" }, { "index": 7, "name": "kv_cache_k", "type": "storage" }, { "index": 8, "name": "kv_cache_v", "type": "storage" } ], "baseUniforms": { "size": 64, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "q_size", "type": "u32", "offset": 4 }, { "name": "k_size", "type": "u32", "offset": 8 }, { "name": "v_size", "type": "u32", "offset": 12 }, { "name": "num_heads", "type": "u32", "offset": 16 }, { "name": "num_kv_heads", "type": "u32", "offset": 20 }, { "name": "head_dim", "type": "u32", "offset": 24 }, { "name": "workgroup_stride", "type": "u32", "offset": 28 }, { "name": "qk_rows", "type": "u32", "offset": 32 }, { "name": "total_v", "type": "u32", "offset": 36 }, { "name": "start_pos", "type": "u32", "offset": 40 }, { "name": "rotary_dim", "type": "u32", "offset": 44 }, { "name": "eps", "type": "f32", "offset": 48 }, { "name": "kv_dst_offset", "type": "u32", "offset": 52 }, { "name": "interleaved", "type": "u32", "offset": 56 }, { "name": "_pad2", "type": "u32", "offset": 60 } ] }, "variants": { "default": { "wgsl": "split_qkv_rmsnorm_rope_qk_f16kv_cache.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "split_qg": { "description": "De-interleave Q and Gate projections from q_proj output for attentionOutputGate models", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "qg_interleaved", "type": "read-only-storage" }, { "index": 2, "name": "Q", "type": "storage" }, { "index": 3, "name": "G", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "num_heads", "type": "u32", "offset": 4 }, { "name": "head_dim", "type": "u32", "offset": 8 }, { "name": "_pad", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "split_qg.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["split-qg.rules.json#variant"], "wgslOnDisk": true } }, "f16": { "wgsl": "split_qg_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["split-qg.rules.json#variant"], "wgslOnDisk": true } } } }, "rep_penalty": { "description": "GPU-side repetition penalty (modifies logits in-place before sampling)", "baseBindings": [ { "index": 0, "type": "uniform" }, { "index": 1, "type": "storage" }, { "index": 2, "type": "read-only-storage" }, { "index": 3, "type": "read-only-storage" } ], "baseUniforms": { "vocab_size": { "type": "u32", "offset": 0 }, "history_count": { "type": "u32", "offset": 4 }, "penalty": { "type": "f32", "offset": 8 }, "batch_count": { "type": "u32", "offset": 12 }, "batch_offset": { "type": "u32", "offset": 16 }, "_pad0": { "type": "u32", "offset": 20 }, "_pad1": { "type": "u32", "offset": 24 }, "_pad2": { "type": "u32", "offset": 28 } }, "variants": { "default": { "wgsl": "rep_penalty.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["rep-penalty.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "rep_penalty_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["rep-penalty.rules.json#variant"], "wgslOnDisk": true } } } }, "sample": { "description": "GPU-side sampling kernels", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "logits", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "vocab_size", "type": "u32", "offset": 0 }, { "name": "temperature", "type": "f32", "offset": 4 }, { "name": "top_k", "type": "u32", "offset": 8 }, { "name": "top_p", "type": "f32", "offset": 12 }, { "name": "random_seed", "type": "u32", "offset": 16 } ] }, "variants": { "argmax": { "wgsl": "sample.wgsl", "entryPoint": "argmax", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "argmax_f16": { "wgsl": "sample_f16.wgsl", "entryPoint": "argmax", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "argmax_reduce": { "wgsl": "sample.wgsl", "entryPoint": "argmax_reduce", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "argmax_reduce_f16": { "wgsl": "sample_f16.wgsl", "entryPoint": "argmax_reduce", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "find_topk_phase1": { "wgsl": "sample.wgsl", "entryPoint": "find_topk_phase1", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "find_topk_phase1_f16": { "wgsl": "sample_f16.wgsl", "entryPoint": "find_topk_phase1", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "find_topk_phase2": { "wgsl": "sample.wgsl", "entryPoint": "find_topk_phase2", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "find_topk_phase2_f16": { "wgsl": "sample_f16.wgsl", "entryPoint": "find_topk_phase2", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "softmax_and_sample": { "wgsl": "sample.wgsl", "entryPoint": "softmax_and_sample", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "softmax_and_sample_f16": { "wgsl": "sample_f16.wgsl", "entryPoint": "softmax_and_sample", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "single_pass": { "wgsl": "sample.wgsl", "entryPoint": "sample_single_pass", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "pinned", "inlineConfigs": [ "diffusiongemma/diffusiongemma-26b-a4b-it-q4k-ehf16-af16", "gemma3/gemma-3-1b-it-f16-af32", "gemma3/gemma-3-1b-it-q4k-ehf16-af32", "gemma3/gemma-3-270m-it-f16-af32", "gemma3/gemma-3-270m-it-q4k-ehf16-af16", "gemma3/gemma-3-270m-it-q4k-ehf16-af32", "gemma3/translategemma-4b-1b-enes-q4k-ehf16-af32", "gemma3/translategemma-4b-it-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-12b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-12b-it-text-q4k-ehf16-hq4k-af16", "gemma4/gemma-4-12b-it-text-w4a16-ct-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af16", "gemma4/gemma-4-31b-it-text-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af16-int4ple", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32", "gemma4/gemma-4-e2b-it-q4k-ehf16-af32-int4ple", "gemma4/gemma-4-moe-q4k-ehf16-af32", "gpt-oss-20b-f16-xmxfp4", "janus/janus-pro-1b-text-q4k-ehaf16", "lfm2/lfm2.5-1.2b-instruct-q4k-ehf16-af32", "minicpm/minicpm4-0-5b-f16-af32", "qwen3/qwen-3-5-0-8b-q4k-ehaf16", "qwen3/qwen-3-5-2b-q4k-ehaf16", "qwen3/qwen-3-5-9b-f16-af32", "qwen3/qwen-3-5-9b-q4k-ehaf16", "qwen3/qwen-3-6-27b-q4k-eaf16", "qwen3/qwen-3-6-27b-q4k-ehaf16", "qwen3/qwen-3-reranker-0-6b-f16-af32", "qwen3/qwen-3-reranker-0-6b-q4k-ehf16-af32" ], "ruleChains": [], "wgslOnDisk": true } }, "single_pass_f16": { "wgsl": "sample_f16.wgsl", "entryPoint": "sample_single_pass", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "bf16_to_f32": { "description": "BF16 to F32 conversion", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 } ] }, "variants": { "default": { "wgsl": "bf16_to_f32.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "bf16_to_f16": { "description": "BF16 to F16 conversion", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 } ] }, "variants": { "default": { "wgsl": "bf16_to_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "transpose": { "description": "Matrix transpose", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "rows", "type": "u32", "offset": 0 }, { "name": "cols", "type": "u32", "offset": 4 }, { "name": "_pad0", "type": "u32", "offset": 8 }, { "name": "_pad1", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "transpose.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "matmul_backward": { "description": "Matmul backward dX = dY @ W^T (fused tiled kernel)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "dY", "type": "read-only-storage" }, { "index": 2, "name": "W", "type": "read-only-storage" }, { "index": 3, "name": "dX", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "M", "type": "u32", "offset": 0 }, { "name": "N", "type": "u32", "offset": 4 }, { "name": "K", "type": "u32", "offset": 8 }, { "name": "alpha", "type": "f32", "offset": 12 }, { "name": "transpose_b", "type": "u32", "offset": 16 } ] }, "variants": { "default": { "wgsl": "backward/matmul_backward.wgsl", "entryPoint": "main", "workgroup": [16,16,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "f16_weight": { "wgsl": "backward/matmul_backward_f16w.wgsl", "entryPoint": "main", "workgroup": [16,16,1], "requires": [], "weightDtype": "f16", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } }, "q4k_weight": { "wgsl": "backward/matmul_backward_q4k.wgsl", "entryPoint": "main", "workgroup": [16,16,1], "requires": [], "weightDtype": "q4k", "description": "Input gradient through a frozen row-wise Q4_K weight stored as [N,K]", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "matmul_transpose_a": { "description": "Matmul with transposed A (C = A^T * B)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "A", "type": "read-only-storage" }, { "index": 2, "name": "B", "type": "read-only-storage" }, { "index": 3, "name": "C", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "M", "type": "u32", "offset": 0 }, { "name": "N", "type": "u32", "offset": 4 }, { "name": "K", "type": "u32", "offset": 8 }, { "name": "alpha", "type": "f32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/matmul_transpose_a.wgsl", "entryPoint": "main", "workgroup": [16,16,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "silu_backward": { "description": "SiLU backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "grad_output", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "_pad0", "type": "u32", "offset": 4 }, { "name": "_pad1", "type": "u32", "offset": 8 }, { "name": "_pad2", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/silu_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "embed_backward": { "description": "Embedding backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "indices", "type": "read-only-storage" }, { "index": 2, "name": "grad_output", "type": "read-only-storage" }, { "index": 3, "name": "grad_embeddings", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "vocab_size", "type": "u32", "offset": 8 }, { "name": "transpose", "type": "u32", "offset": 12 }, { "name": "index_offset", "type": "u32", "offset": 16 }, { "name": "_pad0", "type": "u32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "backward/embed_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "gelu_backward": { "description": "GELU backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "grad_output", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "_pad0", "type": "u32", "offset": 4 }, { "name": "_pad1", "type": "u32", "offset": 8 }, { "name": "_pad2", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/gelu_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "scale_backward": { "description": "Scale backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "grad_output", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "scale", "type": "f32", "offset": 4 }, { "name": "_pad0", "type": "u32", "offset": 8 }, { "name": "_pad1", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/scale_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "softmax_backward": { "description": "Softmax backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "softmax", "type": "read-only-storage" }, { "index": 2, "name": "grad_output", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "rows", "type": "u32", "offset": 0 }, { "name": "cols", "type": "u32", "offset": 4 }, { "name": "_pad0", "type": "u32", "offset": 8 }, { "name": "_pad1", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/softmax_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "attention_backward": { "description": "Attention backward (softmax backward with causal mask)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "softmax", "type": "read-only-storage" }, { "index": 2, "name": "grad_output", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "rows", "type": "u32", "offset": 0 }, { "name": "cols", "type": "u32", "offset": 4 }, { "name": "causal", "type": "u32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/attention_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "rmsnorm_backward": { "description": "RMSNorm backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "grad_output", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "eps", "type": "f32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/rmsnorm_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "gated_rmsnorm_backward": { "description": "Gated RMSNorm backward for frozen norm weights", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "gate", "type": "read-only-storage" }, { "index": 3, "name": "weight", "type": "read-only-storage" }, { "index": 4, "name": "grad_output", "type": "read-only-storage" }, { "index": 5, "name": "grad_input", "type": "storage" }, { "index": 6, "name": "grad_gate", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "rows", "type": "u32", "offset": 0 }, { "name": "width", "type": "u32", "offset": 4 }, { "name": "eps", "type": "f32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/gated_rmsnorm_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "causal_conv1d_silu_backward": { "description": "Causal depthwise Conv1D plus SiLU input backward for frozen weights", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "grad_output", "type": "read-only-storage" }, { "index": 4, "name": "grad_input", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "channels", "type": "u32", "offset": 4 }, { "name": "kernel_size", "type": "u32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/causal_conv1d_silu_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "gated_delta_recurrent_backward": { "description": "Full-history gated-delta recurrent backward oracle", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "query_key", "type": "read-only-storage" }, { "index": 2, "name": "value", "type": "read-only-storage" }, { "index": 3, "name": "decay_beta", "type": "read-only-storage" }, { "index": 4, "name": "state_history", "type": "read-only-storage" }, { "index": 5, "name": "grad_output", "type": "read-only-storage" }, { "index": 6, "name": "grad_query_key", "type": "storage" }, { "index": 7, "name": "grad_value", "type": "storage" }, { "index": 8, "name": "grad_decay_beta", "type": "storage" }, { "index": 9, "name": "grad_state", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "num_heads", "type": "u32", "offset": 4 }, { "name": "key_dim", "type": "u32", "offset": 8 }, { "name": "value_dim", "type": "u32", "offset": 12 }, { "name": "query_scale", "type": "f32", "offset": 16 }, { "name": "_pad0", "type": "u32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "backward/gated_delta_recurrent_backward.wgsl", "entryPoint": "main", "workgroup": [128,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "layernorm_backward": { "description": "LayerNorm backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "grad_output", "type": "read-only-storage" }, { "index": 4, "name": "grad_input", "type": "storage" }, { "index": 5, "name": "grad_weight", "type": "storage" }, { "index": 6, "name": "grad_bias", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "hidden_size", "type": "u32", "offset": 0 }, { "name": "num_tokens", "type": "u32", "offset": 4 }, { "name": "eps", "type": "f32", "offset": 8 }, { "name": "_pad", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/layernorm_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "rope_backward": { "description": "RoPE backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "grad_output", "type": "read-only-storage" }, { "index": 2, "name": "freqs_cos", "type": "read-only-storage" }, { "index": 3, "name": "freqs_sin", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "seq_len", "type": "u32", "offset": 0 }, { "name": "num_heads", "type": "u32", "offset": 4 }, { "name": "head_dim", "type": "u32", "offset": 8 }, { "name": "start_pos", "type": "u32", "offset": 12 }, { "name": "_pad0", "type": "u32", "offset": 16 }, { "name": "_pad1", "type": "u32", "offset": 20 }, { "name": "_pad2", "type": "u32", "offset": 24 }, { "name": "_pad3", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "backward/rope_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "cross_entropy_loss": { "description": "Cross-entropy loss (expects softmax input)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "softmax", "type": "read-only-storage" }, { "index": 2, "name": "targets", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "vocab_size", "type": "u32", "offset": 4 }, { "name": "_pad0", "type": "u32", "offset": 8 }, { "name": "_pad1", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "cross_entropy_loss.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "conv2d": { "description": "Conv2D (NCHW)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "bias", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 48, "fields": [ { "name": "in_channels", "type": "u32", "offset": 0 }, { "name": "out_channels", "type": "u32", "offset": 4 }, { "name": "height", "type": "u32", "offset": 8 }, { "name": "width", "type": "u32", "offset": 12 }, { "name": "out_height", "type": "u32", "offset": 16 }, { "name": "out_width", "type": "u32", "offset": 20 }, { "name": "kernel_h", "type": "u32", "offset": 24 }, { "name": "kernel_w", "type": "u32", "offset": 28 }, { "name": "stride", "type": "u32", "offset": 32 }, { "name": "pad", "type": "u32", "offset": 36 }, { "name": "_pad0", "type": "u32", "offset": 40 }, { "name": "_pad1", "type": "u32", "offset": 44 } ] }, "variants": { "default": { "wgsl": "conv2d.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["conv2d.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "conv2d_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["conv2d.rules.json#variant"], "wgslOnDisk": true } } } }, "depthwise_conv2d": { "description": "Depthwise Conv2D (NCHW)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "bias", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 48, "fields": [ { "name": "channels", "type": "u32", "offset": 0 }, { "name": "height", "type": "u32", "offset": 4 }, { "name": "width", "type": "u32", "offset": 8 }, { "name": "out_height", "type": "u32", "offset": 12 }, { "name": "out_width", "type": "u32", "offset": 16 }, { "name": "kernel_h", "type": "u32", "offset": 20 }, { "name": "kernel_w", "type": "u32", "offset": 24 }, { "name": "stride", "type": "u32", "offset": 28 }, { "name": "pad", "type": "u32", "offset": 32 }, { "name": "_pad0", "type": "u32", "offset": 36 }, { "name": "_pad1", "type": "u32", "offset": 40 }, { "name": "_pad2", "type": "u32", "offset": 44 } ] }, "variants": { "default": { "wgsl": "depthwise_conv2d.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["depthwise-conv2d.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "depthwise_conv2d_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["depthwise-conv2d.rules.json#variant"], "wgslOnDisk": true } } } }, "depthwise_conv1d": { "description": "Depthwise Conv1D (channels-first, same-padded)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "channels", "type": "u32", "offset": 0 }, { "name": "length", "type": "u32", "offset": 4 }, { "name": "kernel_size", "type": "u32", "offset": 8 }, { "name": "_pad0", "type": "u32", "offset": 12 } ] }, "variants": { "depthwise_conv1d": { "wgsl": "depthwise_conv1d.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "grouped_pointwise_conv2d": { "description": "Grouped Pointwise Conv2D (NCHW)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "bias", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "in_channels", "type": "u32", "offset": 0 }, { "name": "out_channels", "type": "u32", "offset": 4 }, { "name": "height", "type": "u32", "offset": 8 }, { "name": "width", "type": "u32", "offset": 12 }, { "name": "groups", "type": "u32", "offset": 16 }, { "name": "_pad0", "type": "u32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "grouped_pointwise_conv2d.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["grouped-pointwise-conv2d.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "grouped_pointwise_conv2d_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["grouped-pointwise-conv2d.rules.json#variant"], "wgslOnDisk": true } } } }, "groupnorm_stats": { "description": "GroupNorm stats (mean/var) for each group", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "stats", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "channels", "type": "u32", "offset": 0 }, { "name": "height", "type": "u32", "offset": 4 }, { "name": "width", "type": "u32", "offset": 8 }, { "name": "num_groups", "type": "u32", "offset": 12 }, { "name": "eps", "type": "f32", "offset": 16 }, { "name": "_pad0", "type": "u32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "stats": { "wgsl": "groupnorm_stats.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["groupnorm.rules.json#statsVariant"], "wgslOnDisk": true } }, "stats_f16": { "wgsl": "groupnorm_stats_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["groupnorm.rules.json#statsVariant"], "wgslOnDisk": true } } } }, "groupnorm_apply": { "description": "GroupNorm apply (NCHW)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "stats", "type": "read-only-storage" }, { "index": 3, "name": "weight", "type": "read-only-storage" }, { "index": 4, "name": "bias", "type": "read-only-storage" }, { "index": 5, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "channels", "type": "u32", "offset": 0 }, { "name": "height", "type": "u32", "offset": 4 }, { "name": "width", "type": "u32", "offset": 8 }, { "name": "num_groups", "type": "u32", "offset": 12 }, { "name": "eps", "type": "f32", "offset": 16 }, { "name": "_pad0", "type": "u32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "apply": { "wgsl": "groupnorm_apply.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["groupnorm.rules.json#applyVariant"], "wgslOnDisk": true } }, "apply_f16": { "wgsl": "groupnorm_apply_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["groupnorm.rules.json#applyVariant"], "wgslOnDisk": true } } } }, "upsample2d": { "description": "Upsample2D (nearest, NCHW)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "channels", "type": "u32", "offset": 0 }, { "name": "in_height", "type": "u32", "offset": 4 }, { "name": "in_width", "type": "u32", "offset": 8 }, { "name": "out_height", "type": "u32", "offset": 12 }, { "name": "out_width", "type": "u32", "offset": 16 }, { "name": "scale", "type": "u32", "offset": 20 }, { "name": "_pad0", "type": "u32", "offset": 24 }, { "name": "_pad1", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "upsample2d.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["upsample2d.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "upsample2d_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["upsample2d.rules.json#variant"], "wgslOnDisk": true } } } }, "modulate": { "description": "Per-channel affine + optional gating", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "mod", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "hidden_size", "type": "u32", "offset": 4 }, { "name": "scale_offset", "type": "u32", "offset": 8 }, { "name": "shift_offset", "type": "u32", "offset": 12 }, { "name": "gate_offset", "type": "u32", "offset": 16 }, { "name": "has_gate", "type": "u32", "offset": 20 }, { "name": "add_one", "type": "u32", "offset": 24 }, { "name": "_pad0", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "modulate.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["modulate.rules.json#variant"], "wgslOnDisk": true } }, "f16": { "wgsl": "modulate_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["modulate.rules.json#variant"], "wgslOnDisk": true } } } }, "pixel_shuffle": { "description": "Pixel shuffle (tokens -> NCHW)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "out_channels", "type": "u32", "offset": 0 }, { "name": "out_height", "type": "u32", "offset": 4 }, { "name": "out_width", "type": "u32", "offset": 8 }, { "name": "grid_width", "type": "u32", "offset": 12 }, { "name": "grid_height", "type": "u32", "offset": 16 }, { "name": "patch_size", "type": "u32", "offset": 20 }, { "name": "patch_channels", "type": "u32", "offset": 24 }, { "name": "_pad0", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "pixel_shuffle.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["pixel_shuffle.rules.json#variant"], "wgslOnDisk": true } }, "f16": { "wgsl": "pixel_shuffle_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["pixel_shuffle.rules.json#variant"], "wgslOnDisk": true } } } }, "cross_entropy_backward": { "description": "Cross-entropy backward (expects softmax input)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "softmax", "type": "read-only-storage" }, { "index": 2, "name": "targets", "type": "read-only-storage" }, { "index": 3, "name": "grad_output", "type": "read-only-storage" }, { "index": 4, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "vocab_size", "type": "u32", "offset": 4 }, { "name": "dispatch_stride", "type": "u32", "offset": 8 }, { "name": "_pad1", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "backward/cross_entropy_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "adam": { "description": "Adam optimizer step", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "params", "type": "storage" }, { "index": 2, "name": "grads", "type": "read-only-storage" }, { "index": 3, "name": "moment1", "type": "storage" }, { "index": 4, "name": "moment2", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "step", "type": "u32", "offset": 4 }, { "name": "lr", "type": "f32", "offset": 8 }, { "name": "beta1", "type": "f32", "offset": 12 }, { "name": "beta2", "type": "f32", "offset": 16 }, { "name": "eps", "type": "f32", "offset": 20 }, { "name": "_pad0", "type": "u32", "offset": 24 }, { "name": "_pad1", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "backward/adam.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "bias_add_backward": { "description": "Bias add backward (sum reduction)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "grad_output", "type": "read-only-storage" }, { "index": 2, "name": "grad_bias", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "num_tokens", "type": "u32", "offset": 0 }, { "name": "dim", "type": "u32", "offset": 4 } ] }, "variants": { "default": { "wgsl": "backward/bias_add_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "upsample2d_backward": { "description": "Upsample2D backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "grad_output", "type": "read-only-storage" }, { "index": 2, "name": "grad_input", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "channels", "type": "u32", "offset": 0 }, { "name": "in_height", "type": "u32", "offset": 4 }, { "name": "in_width", "type": "u32", "offset": 8 }, { "name": "out_height", "type": "u32", "offset": 12 }, { "name": "out_width", "type": "u32", "offset": 16 }, { "name": "scale", "type": "u32", "offset": 20 } ] }, "variants": { "default": { "wgsl": "backward/upsample2d_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "pixel_shuffle_backward": { "description": "Pixel shuffle backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "grad_output", "type": "read-only-storage" }, { "index": 2, "name": "grad_input", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "out_channels", "type": "u32", "offset": 0 }, { "name": "out_height", "type": "u32", "offset": 4 }, { "name": "out_width", "type": "u32", "offset": 8 }, { "name": "grid_width", "type": "u32", "offset": 12 }, { "name": "grid_height", "type": "u32", "offset": 16 }, { "name": "patch_size", "type": "u32", "offset": 20 }, { "name": "patch_channels", "type": "u32", "offset": 24 } ] }, "variants": { "default": { "wgsl": "backward/pixel_shuffle_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "groupnorm_backward": { "description": "GroupNorm backward", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "grad_output", "type": "read-only-storage" }, { "index": 4, "name": "grad_input", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "channels", "type": "u32", "offset": 0 }, { "name": "height", "type": "u32", "offset": 4 }, { "name": "width", "type": "u32", "offset": 8 }, { "name": "num_groups", "type": "u32", "offset": 12 }, { "name": "eps", "type": "f32", "offset": 16 } ] }, "variants": { "default": { "wgsl": "backward/groupnorm_backward.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "linear_attention_summary": { "description": "Linear attention summary stage", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "query", "type": "read-only-storage" }, { "index": 2, "name": "key", "type": "read-only-storage" }, { "index": 3, "name": "value", "type": "read-only-storage" }, { "index": 4, "name": "summary", "type": "storage" } ], "baseUniforms": { "size": 24, "fields": [ { "name": "num_heads", "type": "u32", "offset": 0 }, { "name": "head_dim", "type": "u32", "offset": 4 }, { "name": "num_tokens", "type": "u32", "offset": 8 }, { "name": "hidden_size", "type": "u32", "offset": 12 }, { "name": "_pad0", "type": "u32", "offset": 16 }, { "name": "_pad1", "type": "u32", "offset": 20 } ] }, "variants": { "default": { "wgsl": "linear_attention_summary.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["linear-attention.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "linear_attention_summary_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["linear-attention.rules.json#variant"], "wgslOnDisk": true } } } }, "linear_attention_apply": { "description": "Linear attention apply stage", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "query", "type": "read-only-storage" }, { "index": 2, "name": "summary", "type": "read-only-storage" }, { "index": 3, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 32, "fields": [ { "name": "num_heads", "type": "u32", "offset": 0 }, { "name": "head_dim", "type": "u32", "offset": 4 }, { "name": "num_tokens", "type": "u32", "offset": 8 }, { "name": "hidden_size", "type": "u32", "offset": 12 }, { "name": "eps", "type": "f32", "offset": 16 }, { "name": "_pad0", "type": "u32", "offset": 20 }, { "name": "_pad1", "type": "u32", "offset": 24 }, { "name": "_pad2", "type": "u32", "offset": 28 } ] }, "variants": { "default": { "wgsl": "linear_attention_apply.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["linear-attention.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "linear_attention_apply_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["linear-attention.rules.json#variant"], "wgslOnDisk": true } } } }, "repeat_channels": { "description": "Repeat channels along the NCHW channel axis", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 20, "fields": [ { "name": "in_channels", "type": "u32", "offset": 0 }, { "name": "height", "type": "u32", "offset": 4 }, { "name": "width", "type": "u32", "offset": 8 }, { "name": "repeats", "type": "u32", "offset": 12 }, { "name": "_pad0", "type": "u32", "offset": 16 } ] }, "variants": { "default": { "wgsl": "repeat_channels.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["repeat-channels.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "repeat_channels_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["repeat-channels.rules.json#variant"], "wgslOnDisk": true } } } }, "relu": { "description": "ReLU activation", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "input", "type": "read-only-storage" }, { "index": 2, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 16, "fields": [ { "name": "size", "type": "u32", "offset": 0 }, { "name": "_pad0", "type": "u32", "offset": 4 }, { "name": "_pad1", "type": "u32", "offset": 8 }, { "name": "_pad2", "type": "u32", "offset": 12 } ] }, "variants": { "default": { "wgsl": "relu.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "outputDtype": "f32", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["relu.rules.json#variant"], "wgslOnDisk": true } }, "default_f16": { "wgsl": "relu_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "reachability": { "status": "selectable", "inlineConfigs": [], "ruleChains": ["relu.rules.json#variant"], "wgslOnDisk": true } } } }, "conv2d_backward_input": { "description": "Conv2D backward (input gradient)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "grad_output", "type": "read-only-storage" }, { "index": 2, "name": "weight", "type": "read-only-storage" }, { "index": 3, "name": "grad_input", "type": "storage" } ], "baseUniforms": { "size": 48, "fields": [ { "name": "in_channels", "type": "u32", "offset": 0 }, { "name": "out_channels", "type": "u32", "offset": 4 }, { "name": "height", "type": "u32", "offset": 8 }, { "name": "width", "type": "u32", "offset": 12 }, { "name": "out_height", "type": "u32", "offset": 16 }, { "name": "out_width", "type": "u32", "offset": 20 }, { "name": "kernel_h", "type": "u32", "offset": 24 }, { "name": "kernel_w", "type": "u32", "offset": 28 }, { "name": "stride", "type": "u32", "offset": 32 }, { "name": "pad", "type": "u32", "offset": 36 } ] }, "variants": { "default": { "wgsl": "backward/conv2d_backward_input.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "conv2d_backward_weight": { "description": "Conv2D backward (weight gradient)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "grad_output", "type": "read-only-storage" }, { "index": 2, "name": "input", "type": "read-only-storage" }, { "index": 3, "name": "grad_weight", "type": "storage" } ], "baseUniforms": { "size": 48, "fields": [ { "name": "in_channels", "type": "u32", "offset": 0 }, { "name": "out_channels", "type": "u32", "offset": 4 }, { "name": "height", "type": "u32", "offset": 8 }, { "name": "width", "type": "u32", "offset": 12 }, { "name": "out_height", "type": "u32", "offset": 16 }, { "name": "out_width", "type": "u32", "offset": 20 }, { "name": "kernel_h", "type": "u32", "offset": 24 }, { "name": "kernel_w", "type": "u32", "offset": 28 }, { "name": "stride", "type": "u32", "offset": 32 }, { "name": "pad", "type": "u32", "offset": 36 } ] }, "variants": { "default": { "wgsl": "backward/conv2d_backward_weight.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": [], "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } }, "attention_bdpa": { "description": "Basis-decomposed paged attention kernels (decode)", "baseBindings": [ { "index": 0, "name": "uniforms", "type": "uniform" }, { "index": 1, "name": "Q", "type": "read-only-storage" }, { "index": 2, "name": "T_basis_k", "type": "read-only-storage" }, { "index": 3, "name": "T_basis_v", "type": "read-only-storage" }, { "index": 4, "name": "P_delta_k", "type": "read-only-storage" }, { "index": 5, "name": "P_delta_v", "type": "read-only-storage" }, { "index": 6, "name": "I_flat", "type": "read-only-storage" }, { "index": 7, "name": "rope_cos", "type": "read-only-storage" }, { "index": 8, "name": "rope_sin", "type": "read-only-storage" }, { "index": 9, "name": "output", "type": "storage" } ], "baseUniforms": { "size": 64, "fields": [ { "name": "num_heads", "type": "u32", "offset": 0 }, { "name": "num_kv_heads", "type": "u32", "offset": 4 }, { "name": "head_dim", "type": "u32", "offset": 8 }, { "name": "kv_len", "type": "u32", "offset": 12 }, { "name": "seq_len", "type": "u32", "offset": 16 }, { "name": "scale", "type": "f32", "offset": 20 }, { "name": "is_causal", "type": "u32", "offset": 24 }, { "name": "start_pos", "type": "u32", "offset": 28 }, { "name": "attn_softcap", "type": "f32", "offset": 32 }, { "name": "sliding_window", "type": "u32", "offset": 36 }, { "name": "kv_len_source", "type": "u32", "offset": 40 }, { "name": "_pad_1", "type": "u32", "offset": 44 }, { "name": "_pad_2", "type": "u32", "offset": 48 }, { "name": "_pad_3", "type": "u32", "offset": 52 }, { "name": "_pad_4", "type": "u32", "offset": 56 }, { "name": "_pad_5", "type": "u32", "offset": 60 } ] }, "variants": { "decode_bdpa_f16": { "wgsl": "attention_bdpa_decode_f16.wgsl", "entryPoint": "main", "workgroup": [256,1,1], "requires": ["shader-f16"], "outputDtype": "f16", "variantMetadata": { "maxKVLen": 2048 }, "description": "BDPA decode kernel - f16 Q/basis + packed int8 deltas", "reachability": { "status": "unused", "inlineConfigs": [], "ruleChains": [], "wgslOnDisk": true } } } } } }