Bump hidden dimension size for B300 kernels with Float8Block.

Signed-off-by: Cory Ye <cye@nvidia.com>

Author: cspades

tests/pytorch/distributed/run_fsdp2_allgather.py

@@ -32,9 +32,9 @@
 
 # Fixed model dimensions — this test focuses on allgather correctness, not model flexibility.
 _NUM_HEADS = 8
-_HEAD_DIM = 64
-_HIDDEN_SIZE = _NUM_HEADS * _HEAD_DIM  # 512
-_FFN_SIZE = _HIDDEN_SIZE * 4  # 2048
+_HEAD_DIM = 128
+_HIDDEN_SIZE = _NUM_HEADS * _HEAD_DIM
+_FFN_SIZE = _HIDDEN_SIZE * 4
 _NUM_LAYERS = 2
 _BATCH_SIZE = 4
 _SEQ_LEN = 32

tests/pytorch/distributed/run_fsdp2_model.py

@@ -102,7 +102,7 @@ def dist_print(msg):
 def _parse_args(argv=None, namespace=None):
     parser = argparse.ArgumentParser(description="Toy example for debugging fully_shard()")
     parser.add_argument("--num-heads", type=int, default=8, help="Number of attn. heads")
-    parser.add_argument("--head-dim", type=int, default=64, help="Attention head size")
+    parser.add_argument("--head-dim", type=int, default=128, help="Attention head size")
     parser.add_argument("--batch-size", type=int, default=16, help="Batch size of input")
     parser.add_argument("--seq-length", type=int, default=128, help="Sequence length of input")
     parser.add_argument("--params-dtype", type=str, default="float32", help="Parameter dtype.")

transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py

@@ -63,7 +63,7 @@ def __new__(
         instance._rowwise_scale_inv = rowwise_scale_inv
         instance._columnwise_scale_inv = columnwise_scale_inv
         instance._is_2D_scaled = is_2D_scaled
-        instance._default_storage = torch.UntypedStorage(1)
+        instance._default_storage = torch.UntypedStorage(1, device=torch.cuda.current_device())
 
         return instance
 

transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py

@@ -101,7 +101,7 @@ def __new__(
         instance._quantizer = quantizer.copy() if quantizer is not None else None
         instance._fp8_dtype = fp8_dtype
         instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales
-        instance._default_storage = torch.UntypedStorage(1)
+        instance._default_storage = torch.UntypedStorage(1, device=torch.cuda.current_device())
 
         return instance
 

transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py

@@ -126,7 +126,7 @@ def __new__(
         instance._amax_rowwise = amax_rowwise
         instance._amax_columnwise = amax_columnwise
         instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales
-        instance._default_storage = torch.UntypedStorage(1)
+        instance._default_storage = torch.UntypedStorage(1, device=torch.cuda.current_device())
 
         return instance