Preserve Tenstorrent shim compatibility

Author: Andrey Cheptsov

mkdocs/docs/examples/accelerators/tenstorrent.md

@@ -201,16 +201,17 @@ If you run it via `dstack apply`, it will output the URL to access it via your d
 
 `resources.gpu` uses the usual `name:count` format. For Tenstorrent, `count`
 is the number of devices reported from the TT-SMI topology. On Galaxy systems,
-this corresponds to chips. For PCIe cards and workstations, use the card count
-shown in the examples below.
+this corresponds to chips. On PCIe systems, this is usually the card count, but
+dual-chip cards can also be reported as per-chip devices.
 
 ```yaml
 resources:
   gpu: tt-galaxy-wh:32  # Galaxy Wormhole, 32 chips
   # gpu: tt-galaxy-bh:32  # Galaxy Blackhole, 32 chips
   # gpu: n300:4           # TT-LoudBox or TT-QuietBox Wormhole, 4 n300 cards
   # gpu: p150:4           # TT-QuietBox Blackhole, 4 p150 cards
-  # gpu: p300:2           # TT-QuietBox 2 Blackhole, 2 p300 cards
+  # gpu: p300:64GB:2      # TT-QuietBox 2 Blackhole, 2 p300 cards
+  # gpu: p300:32GB:4      # TT-QuietBox 2 Blackhole, if exposed per chip
 ```
 
 Use `tt:<count>` only when the workload can run on any Tenstorrent device type.

runner/internal/shim/host/gpu.go

@@ -248,9 +248,13 @@ func splitTtBoardType(boardType string) (name string, suffix string) {
 }
 
 func ttBoardVramMib(name string) int {
-	switch name {
-	case "n150", "n300", "tt-galaxy-wh":
+	switch {
+	case strings.HasPrefix(name, "n150"),
+		strings.HasPrefix(name, "n300"),
+		strings.HasPrefix(name, "tt-galaxy-wh"):
 		return 12 * 1024
+	}
+	switch name {
 	case "p100a":
 		return 28 * 1024
 	case "p150", "p300", "tt-galaxy-bh":
@@ -317,7 +321,7 @@ func getGpusFromTtSmiSnapshot(snapshot *ttSmiSnapshot) []GpuInfo {
 			// and add memory to the first "L" device we find with that board_id
 			for _, key := range gpuKeys {
 				gpu := gpuMap[key]
-				if gpu.ID == boardID && gpu.Name == name {
+				if gpu.ID == boardID && (gpu.Name == name || !isTtBlackholeBoard(name)) {
 					// Add memory to the "L" device
 					gpu.Vram += ttBoardVramMib(name)
 					break // Only add to the first matching "L" device
@@ -358,7 +362,7 @@ func getGpusFromTtSmiSnapshot(snapshot *ttSmiSnapshot) []GpuInfo {
 			existingGpu := false
 			for _, key := range gpuKeys {
 				gpu := gpuMap[key]
-				if gpu.ID == boardID && gpu.Name == name {
+				if gpu.ID == boardID {
 					gpu.Vram += baseVram
 					existingGpu = true
 					break

runner/internal/shim/host/gpu_test.go

@@ -239,6 +239,28 @@ func TestGetGpusFromTtSmiSnapshotMultipleDevices(t *testing.T) {
 	}
 }
 
+func TestGetGpusFromTtSmiSnapshotWormholePrefixMemoryCompatibility(t *testing.T) {
+	snapshot := &ttSmiSnapshot{
+		DeviceInfo: []ttDeviceInfo{
+			{BoardInfo: ttBoardInfo{BoardType: "n150-custom L", BoardID: "100018000000001"}},
+			{BoardInfo: ttBoardInfo{BoardType: "n300-custom L", BoardID: "100014000000001"}},
+			{BoardInfo: ttBoardInfo{BoardType: "n300-custom R", BoardID: "100014000000001"}},
+			{BoardInfo: ttBoardInfo{BoardType: "tt-galaxy-wh-custom L", BoardID: "100035000000001"}},
+		},
+	}
+
+	gpus := getGpusFromTtSmiSnapshot(snapshot)
+
+	expected := []GpuInfo{
+		{Vendor: gpu.GpuVendorTenstorrent, Name: "n150-custom", Vram: 12 * 1024, ID: "100018000000001", Index: "0"},
+		{Vendor: gpu.GpuVendorTenstorrent, Name: "n300-custom", Vram: 24 * 1024, ID: "100014000000001", Index: "1"},
+		{Vendor: gpu.GpuVendorTenstorrent, Name: "tt-galaxy-wh-custom", Vram: 12 * 1024, ID: "100035000000001", Index: "2"},
+	}
+	if !reflect.DeepEqual(gpus, expected) {
+		t.Errorf("getGpusFromTtSmiSnapshot() = %v, want %v", gpus, expected)
+	}
+}
+
 func TestGetGpusFromTtSmiSnapshotGalaxy(t *testing.T) {
 	data, err := loadTestData("tenstorrent/galaxy.json")
 	if err != nil {
@@ -329,7 +351,8 @@ func TestGetGpusFromTtSmiSnapshotBlackholeRevisions(t *testing.T) {
 }
 
 func TestGetGpusFromTtSmiSnapshotBlackholeSourceFixtures(t *testing.T) {
-	// Derived from TT-Metal UMD Blackhole board descriptors.
+	// Synthetic tt-smi snapshot derived from TT-SMI board name mappings and
+	// TT-Metal UMD Blackhole board descriptors.
 	data, err := loadTestData("tenstorrent/blackhole_boards.json")
 	if err != nil {
 		t.Fatalf("Failed to load test data: %v", err)
@@ -355,6 +378,7 @@ func TestGetGpusFromTtSmiSnapshotBlackholeSourceFixtures(t *testing.T) {
 
 func TestGetGpusFromTtSmiSnapshotBlackholeEightP150(t *testing.T) {
 	// Derived from TT-Metal UMD blackhole_8xP150 cluster descriptor.
+	// The p150b name follows TT-SMI's board ID to board type mapping.
 	data, err := loadTestData("tenstorrent/blackhole_8xp150.json")
 	if err != nil {
 		t.Fatalf("Failed to load test data: %v", err)

runner/internal/shim/host/testdata/tenstorrent/README.md

@@ -0,0 +1,19 @@
+# Tenstorrent fixtures
+
+The Wormhole fixtures are captured `tt-smi -s` snapshots used by the existing
+Tenstorrent tests.
+
+The Blackhole fixtures are source-derived compatibility fixtures:
+
+- `blackhole_boards.json` covers `tt-smi` Blackhole board names and P300
+  same-board dual-MMIO behavior. Board names are based on
+  `tt_smi/utils.py::get_board_type` and UMD board type mappings.
+- `blackhole_8xp150.json` is derived from UMD's
+  `blackhole_8xP150.yaml` cluster descriptor. The board IDs and PCI bus IDs are
+  from that descriptor; the `p150b` board name follows the UPI mapping used by
+  `tt-smi`.
+- `blackhole_galaxy.json` is derived from the Blackhole Galaxy example in the
+  `tt-smi` README, which shows a 32-ASIC Galaxy reporting `tt-galaxy-bh`.
+
+These fixtures are not substitutes for live hardware smoke tests. They preserve
+the `tt-smi` JSON shapes and UMD topology cases that we can verify from source.