[Feat] Expose attn, batch, ubatch, cach_type_kv settings to the UI an…

…d bench results (#148) * feat: expose attn, cache_k and cache_v types, n_batch, n_ubatch in the ui * feat: add model init args to the benchmark
a-ghorbani · Dec 26, 2024 · b836d07 · b836d07
1 parent f655ef3
commit b836d07
Show file tree

Hide file tree

Showing 8 changed files with 662 additions and 146 deletions.
diff --git a/src/screens/BenchmarkScreen/BenchResultCard/BenchResultCard.tsx b/src/screens/BenchmarkScreen/BenchResultCard/BenchResultCard.tsx
@@ -79,12 +79,42 @@ export const BenchResultCard = ({result, onDelete, onShare}: Props) => {
           </Button>
         </View>
 
-        <View style={styles.configBar}>
-          <Text variant="labelSmall">Config</Text>
-          <Text style={styles.configText}>
-            PP: {result.config.pp} • TG: {result.config.tg} • PL:{' '}
-            {result.config.pl} • Rep: {result.config.nr}
-          </Text>
+        <View style={styles.configContainer}>
+          <View style={styles.configBar}>
+            <Text variant="labelSmall">Benchmark Config</Text>
+            <Text style={styles.configText}>
+              PP: {result.config.pp} • TG: {result.config.tg} • PL:{' '}
+              {result.config.pl} • Rep: {result.config.nr}
+            </Text>
+          </View>
+
+          {result.initSettings && (
+            <View style={styles.configBar}>
+              <Text variant="labelSmall">Model Settings</Text>
+              <View style={styles.configTextContainer}>
+                <Text style={styles.configText}>
+                  Context: {result.initSettings.n_context} • Batch:{' '}
+                  {result.initSettings.n_batch} • UBatch:{' '}
+                  {result.initSettings.n_ubatch}
+                </Text>
+                <Text style={styles.configText}>
+                  CPU Threads: {result.initSettings.n_threads} • GPU Layers:{' '}
+                  {result.initSettings.n_gpu_layers}
+                </Text>
+                {result.initSettings.flash_attn ? (
+                  <Text style={styles.configText}>
+                    Flash Attention Enabled • Cache Types:{' '}
+                    {result.initSettings.cache_type_k}/
+                    {result.initSettings.cache_type_v}
+                  </Text>
+                ) : (
+                  <Text style={styles.configText}>
+                    Flash Attention Disabled
+                  </Text>
+                )}
+              </View>
+            </View>
+          )}
         </View>
 
         <View style={styles.resultsContainer}>

diff --git a/src/screens/BenchmarkScreen/BenchResultCard/styles.ts b/src/screens/BenchmarkScreen/BenchResultCard/styles.ts
@@ -29,20 +29,27 @@ export const createStyles = (theme: Theme) =>
       fontSize: 12,
       color: theme.colors.onSurfaceVariant,
     },
-    configBar: {
-      flexDirection: 'row',
-      justifyContent: 'space-between',
-      alignItems: 'center',
-      paddingVertical: 8,
-      marginBottom: 16,
+    configContainer: {
       borderTopWidth: 1,
       borderBottomWidth: 1,
       borderColor: theme.colors.surfaceVariant,
+      marginVertical: 8,
+      paddingHorizontal: 12,
+    },
+    configBar: {
+      flexDirection: 'column',
+      alignItems: 'flex-start',
+      paddingVertical: 8,
+      gap: 4,
     },
     configText: {
       fontSize: 12,
       color: theme.colors.onSurfaceVariant,
-      textAlign: 'center',
+      flex: 1,
+    },
+    configTextContainer: {
+      gap: 4,
+      width: '100%',
     },
     resultsContainer: {
       marginBottom: 16,

diff --git a/src/screens/BenchmarkScreen/BenchmarkScreen.tsx b/src/screens/BenchmarkScreen/BenchmarkScreen.tsx
@@ -30,11 +30,11 @@ const DEFAULT_CONFIGS: BenchmarkConfig[] = [
 
 const BENCHMARK_PARAMS_METADATA = {
   pp: {
-    validation: {min: 64, max: 512},
+    validation: {min: 64, max: 2048},
     descriptionKey: 'Number of prompt processing tokens',
   },
   tg: {
-    validation: {min: 32, max: 512},
+    validation: {min: 32, max: 2048},
     descriptionKey: 'Number of text generation tokens',
   },
   pl: {
@@ -176,6 +176,7 @@ export const BenchmarkScreen: React.FC = observer(() => {
         peakMemoryUsage: peakMemoryUsage || undefined,
         wallTimeMs,
         uuid: uuidv4(),
+        initSettings: modelStore.activeContextSettings,
       };
 
       benchmarkStore.addResult(result);