diff --git a/flash.h b/flash.h
index 027392f9e..506dd9db1 100644
--- a/flash.h
+++ b/flash.h
@@ -149,8 +149,10 @@ enum chipbustype {
  */
 #define NUM_ERASEFUNCTIONS 5
 
-#define FEATURE_REGISTERMAP (1 << 0)
-#define FEATURE_BYTEWRITES (1 << 1)
+#define FEATURE_REGISTERMAP	(1 << 0)
+#define FEATURE_BYTEWRITES	(1 << 1)
+#define FEATURE_ADDR_FULL	(0 << 2)
+#define FEATURE_ADDR_MASK	(3 << 2)
 
 struct flashchip {
 	const char *vendor;
diff --git a/jedec.c b/jedec.c
index 3a2153930..055fbc37d 100644
--- a/jedec.c
+++ b/jedec.c
@@ -373,6 +373,19 @@ retry:
 	return failed;
 }
 
+int getaddrmask(struct flashchip *flash)
+{
+	switch (flash->feature_bits & FEATURE_ADDR_MASK) {
+	case FEATURE_ADDR_FULL:
+		return MASK_FULL;
+		break;
+	default:
+		fprintf(stderr, "%s called with unknown mask\n", __func__);
+		return 0;
+		break;
+	}
+}
+
 int write_jedec(struct flashchip *flash, uint8_t *buf)
 {
 	int i, failed = 0;
@@ -438,7 +451,10 @@ int erase_chip_block_jedec(struct flashchip *flash, unsigned int addr,
 
 int probe_jedec(struct flashchip *flash)
 {
-	return probe_jedec_common(flash, MASK_FULL, 1);
+	int mask;
+
+	mask = getaddrmask(flash);
+	return probe_jedec_common(flash, mask, 1);
 }
 
 int erase_sector_jedec(struct flashchip *flash, unsigned int page, unsigned int size)