Skip to content

Commit bf74abf

Browse files
Lizhe Jifacebook-github-bot
authored andcommitted
Add DRAM_SSD backend type to BackendType enum
Summary: X-link: facebookresearch/FBGEMM#2730 Add `DRAM_SSD` as a new backend type to the fbgemm `BackendType` enum in `fbgemm_gpu/tbe/ssd/ssd_config.py`. `DRAM_SSD` represents a composite DRAM + SSD KV backend for KVZCH virtual tables and is required by the `DRAM_SSD_VIRTUAL_TABLE` compute kernel, which references `BackendType.DRAM_SSD`. - Add `DRAM_SSD = 3` to `BackendType`. - Add `"dram_ssd"` mapping to `BackendType.from_str`. Differential Revision: D106733255
1 parent 1509423 commit bf74abf

2 files changed

Lines changed: 13 additions & 0 deletions

File tree

fbgemm_gpu/fbgemm_gpu/tbe/ssd/ssd_config.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,13 +314,15 @@ class BackendType(enum.IntEnum):
314314
SSD = 0
315315
DRAM = 1
316316
PS = 2
317+
DRAM_SSD = 3
317318

318319
@classmethod
319320
def from_str(cls, key: str) -> "BackendType":
320321
lookup = {
321322
"ssd": BackendType.SSD,
322323
"dram": BackendType.DRAM,
323324
"ps": BackendType.PS,
325+
"dram_ssd": BackendType.DRAM_SSD,
324326
}
325327
if key in lookup:
326328
return lookup[key]

fbgemm_gpu/test/tbe/ssd/ssd_config_test.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,17 @@ def test_dram(self) -> None:
123123

124124
self.assertEqual(BackendType.from_str("dram"), BackendType.DRAM)
125125

126+
def test_ps(self) -> None:
127+
from fbgemm_gpu.tbe.ssd import BackendType
128+
129+
self.assertEqual(BackendType.from_str("ps"), BackendType.PS)
130+
131+
def test_dram_ssd(self) -> None:
132+
from fbgemm_gpu.tbe.ssd import BackendType
133+
134+
self.assertEqual(BackendType.from_str("dram_ssd"), BackendType.DRAM_SSD)
135+
self.assertEqual(BackendType.DRAM_SSD.value, 3)
136+
126137
def test_invalid_raises(self) -> None:
127138
from fbgemm_gpu.tbe.ssd import BackendType
128139

0 commit comments

Comments
 (0)