diff --git a/.vscode/launch.json b/.vscode/launch.json
index 83822569..d6f08f5e 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -166,7 +166,7 @@
             "program": "${workspaceFolder}/samples/fir",
             "args": [
                 "-timing",
-                "-length=8192",
+                "-length=81920",
                 "-report-all",
             ],
         },
diff --git a/driver/builder.go b/driver/builder.go
index 3cf47ed8..62f62060 100644
--- a/driver/builder.go
+++ b/driver/builder.go
@@ -17,6 +17,7 @@ type Builder struct {
 	useMagicMemoryCopy  bool
 	middlewareD2HCycles int
 	middlewareH2DCycles int
+	memorySize          uint64
 }
 
 // MakeBuilder creates a driver builder with some default configuration
@@ -74,6 +75,12 @@ func (b Builder) WithH2DCycles(h2dCycles int) Builder {
 	return b
 }
 
+// WithMemorySize sets the memory size of the CPU.
+func (b Builder) WithMemorySize(memorySize uint64) Builder {
+	b.memorySize = memorySize
+	return b
+}
+
 // Build creates a driver.
 func (b Builder) Build(name string) *Driver {
 	driver := new(Driver)
@@ -125,7 +132,7 @@ func (b *Builder) createCPU(d *Driver) {
 		Type:     internal.DeviceTypeCPU,
 		MemState: internal.NewDeviceMemoryState(d.Log2PageSize),
 	}
-	cpu.SetTotalMemSize(4 * mem.GB)
+	cpu.SetTotalMemSize(b.memorySize)
 
 	d.memAllocator.RegisterDevice(cpu)
 	d.devices = append(d.devices, cpu)
diff --git a/driver/driver.go b/driver/driver.go
index 28d17459..bb882d9c 100644
--- a/driver/driver.go
+++ b/driver/driver.go
@@ -1,6 +1,7 @@
 package driver
 
 import (
+	"fmt"
 	"log"
 	"reflect"
 	"runtime/debug"
@@ -455,6 +456,9 @@ func (d *Driver) distributeWGToGPUs(
 		panic("not all wg allocated")
 	}
 
+	// fmt.Sprintln("total WG: %d WG Per CU %d\n", totalWGCount, wgPerCU)
+	fmt.Printf("total WG: %d WG Per CU %d\n", totalWGCount, wgPerCU)
+
 	return wgDist
 }
 
diff --git a/timing/cp/builder.go b/timing/cp/builder.go
index 797a5bfb..99e65e1e 100644
--- a/timing/cp/builder.go
+++ b/timing/cp/builder.go
@@ -86,9 +86,9 @@ func (b Builder) Build(name string) *CommandProcessor {
 		tracing.CollectTrace(cp, b.visTracer)
 	}
 
-	if b.perfAnalyzer != nil {
-		b.perfAnalyzer.RegisterComponent(cp)
-	}
+	// if b.perfAnalyzer != nil {
+	// 	b.perfAnalyzer.RegisterComponent(cp)
+	// }
 
 	return cp
 }