[Mosaic TPU] No longer implicitly trigger single-buffering when running out of VMEM

It's just more predictable to OOM and let the user explicitly opt into single buffering
using `pipeline_mode=pl.Buffered(1)`. We bump the version number so that the old artifacts
still get the relaxed treatment.

PiperOrigin-RevId: 820593493

Author: apaszke

jax/_src/tpu_custom_call.py

@@ -58,20 +58,20 @@
 # mode: for 1 month when exporting, or when using old cloud TPU.
 #
 # This can be achieved by adding:
-#    if ctx.is_forward_compat() or is_cloud_tpu_older_than(<today>):
+#    if ctx.is_forward_compat() or backend is None or is_cloud_tpu_older_than(<today>):
 #       return <previous_serialization_version>
 #    return None
 #
 # We should also add a TODO to remove the conditional one month later.
 def get_ir_version(ctx: mlir.LoweringRuleContext) -> int | None:
   backend = ctx.module_context.get_backend(optional=True)
-  # TODO(naumsmogers): remove the forward compatibility check after 2025-09-14.
+  # TODO(apaszke): remove the forward compatibility check after 2025-12-5.
   if (
       ctx.is_forward_compat()
       or backend is None
-      or is_cloud_tpu_older_than(2025, 8, 14, backend)
+      or is_cloud_tpu_older_than(2025, 11, 5, backend)
   ):
-    return 7
+    return 8
   return None
 
 

jaxlib/mosaic/dialect/tpu/transforms/serde.cc

@@ -40,7 +40,7 @@ constexpr StringRef kMangledDialect = "stable_mosaic.";
 constexpr StringRef kVersionAttrName = "stable_mosaic.version";
 // When this is bumped, we should file a TODO to update the forward-compatible
 // version in tpu_custom_call.py in a month!
-constexpr int kVersion = 8;
+constexpr int kVersion = 9;
 
 using SerdeRuleType = jaxlib::mosaic::SerdeRuleType;
 
@@ -311,7 +311,8 @@ void MosaicSerdePass::runOnOperation() {
           {.dialect_prefix = kMangledDialect,
            .highest_version = kVersion,
            .version_attr_name = kVersionAttrName,
-           .serialize_version = serialize_version}))) {
+           .serialize_version = serialize_version},
+          /*keep_version_attr=*/keep_version_attr))) {
     signalPassFailure();
   }
 }

jaxlib/mosaic/dialect/tpu/transforms/serde.h

@@ -62,6 +62,8 @@ struct MosaicSerdePass : public jaxlib::mlir::Pass<MosaicSerdePass, ModuleOp> {
 
  protected:
   ::mlir::Pass::Option<bool> serialize{*this, "serialize", llvm::cl::desc("")};
+  ::mlir::Pass::Option<bool> keep_version_attr{
+      *this, "keep-version-attr", llvm::cl::desc(""), llvm::cl::init(true)};
   ::mlir::Pass::Option<int> target_version{*this, "target-version",
                                            llvm::cl::desc("")};
 };

jaxlib/mosaic/serde.cc

@@ -55,7 +55,7 @@ std::optional<llvm::StringRef> demangle(llvm::StringRef name,
 mlir::LogicalResult RunSerde(
     mlir::ModuleOp module, const llvm::StringMap<SerdeRuleType>& upgrade_rules,
     const llvm::StringMap<SerdeRuleType>& downgrade_rules, bool serialize,
-    SerdeOptions options) {
+    SerdeOptions options, bool keep_version_attr) {
   int version = options.highest_version;
   int serialize_version = options.serialize_version;
   if (!serialize && serialize_version != -1) {
@@ -91,7 +91,9 @@ mlir::LogicalResult RunSerde(
       return mlir::failure();
     }
     version = version_attr.getInt();
-    module->removeAttr(options.version_attr_name);
+    if (!keep_version_attr) {
+      module->removeAttr(options.version_attr_name);
+    }
   }
   std::string storage;
   // Explicitly use a post-order walk to allow for deleting operations on the

jaxlib/mosaic/serde.h

@@ -49,9 +49,9 @@ using SerdeRuleType =
 // Run serialization or deserialization on the given module.
 ::mlir::LogicalResult RunSerde(
     ::mlir::ModuleOp module,
-    const llvm::StringMap<SerdeRuleType> &upgrade_rules,
-    const llvm::StringMap<SerdeRuleType> &downgrade_rules, bool serialize,
-    SerdeOptions options);
+    const llvm::StringMap<SerdeRuleType>& upgrade_rules,
+    const llvm::StringMap<SerdeRuleType>& downgrade_rules, bool serialize,
+    SerdeOptions options, bool keep_version_attr = false);
 
 }  // namespace jaxlib::mosaic
 

tests/pallas/tpu_pallas_test.py

@@ -2101,7 +2101,7 @@ def kernel(x_ref, y_ref):
       pl.Buffered(1),
       pl.Buffered(2),
   ])
-  def test_vmem_oom_error_message_basics(self, pmode):
+  def test_vmem_oom_error_message_basics(self, pmode: pl.Buffered):
     if not jtu.if_cloud_tpu_at_least(2025, 10, 14):
       self.skipTest('Support added on Oct 14, 2025')
 
@@ -2152,11 +2152,13 @@ def index_map(i, j):
         f' full shape is f32[{shape[0]},{shape[1]}].',
         error_message,
     )
-    # When VMEM is OOM, double buffering is disabled.
-    self.assertIn(
-        'This allocation is single buffered.',
-        error_message,
-    )
+    if jtu.if_cloud_tpu_at_least(2025, 11, 5):
+      self.assertIn(
+          'This allocation is single buffered.'
+          if pmode.buffer_count == 1
+          else 'This allocation has 2 buffering levels',
+          error_message,
+      )
 
   def test_vmem_oom_error_message_dynamic_grid_scalar_prefetch_and_vmem_scratch(
       self,