@@ -349,6 +349,42 @@ def convert_to_tokamax_splash_config(
349349 )
350350
351351
352+ def _extract_custom_block_sizes (flash_block_sizes ):
353+ """Pulls custom-kernel block sizes out of the (dict or BlockSizes-like) config.
354+
355+ Mirrors the extraction used by the `ulysses_custom` path so the custom ring
356+ kernel honors the same `flash_block_sizes={...}` knobs.
357+ """
358+ bq = 4864
359+ bkv = 1024
360+ bkv_compute = 1024
361+ bkv_compute_in = 1024
362+ heads_per_tile = 1
363+ vmem_limit_bytes = None
364+ if flash_block_sizes is not None :
365+ if isinstance (flash_block_sizes , dict ):
366+ get = flash_block_sizes .get
367+ bq = get ("block_q" , None ) or bq
368+ bkv = get ("block_kv" , None ) or bkv
369+ bkv_compute = get ("block_kv_compute" , None ) or bkv_compute
370+ bkv_compute_in = get ("block_kv_compute_in" , None ) or bkv_compute_in
371+ heads_per_tile = get ("heads_per_tile" , None ) or heads_per_tile
372+ vmem_limit_bytes = get ("vmem_limit_bytes" , None ) or vmem_limit_bytes
373+ else :
374+ bq = getattr (flash_block_sizes , "block_q" , None ) or bq
375+ bkv = getattr (flash_block_sizes , "block_kv" , None ) or bkv
376+ bkv_compute = getattr (flash_block_sizes , "block_kv_compute" , None ) or bkv_compute
377+ bkv_compute_in = getattr (flash_block_sizes , "block_kv_compute_in" , None ) or bkv_compute_in
378+ heads_per_tile = getattr (flash_block_sizes , "heads_per_tile" , None ) or heads_per_tile
379+ vmem_limit_bytes = getattr (flash_block_sizes , "vmem_limit_bytes" , None ) or vmem_limit_bytes
380+ # A BlockSizes object carries heads_per_tile=None when the config dict omitted
381+ # it; getattr then returns that None instead of the default, so coerce it back
382+ # to 1 (the custom-kernel default) to keep the `heads_per_tile > 1` guards safe.
383+ if heads_per_tile is None :
384+ heads_per_tile = 1
385+ return bq , bkv , bkv_compute , bkv_compute_in , heads_per_tile , vmem_limit_bytes
386+
387+
352388def _build_padding_segment_ids (
353389 query_seq_len : int ,
354390 q_padded_len : int ,
@@ -885,6 +921,181 @@ def wrap_ulysses_ring_attention(query, key, value):
885921 return x
886922
887923
924+ def _ulysses_ring_custom_attention (
925+ query : jax .Array ,
926+ key : jax .Array ,
927+ value : jax .Array ,
928+ heads : int ,
929+ mesh : Mesh ,
930+ axis_names_q : AxisNames ,
931+ axis_names_kv : AxisNames ,
932+ flash_block_sizes : BlockSizes ,
933+ dtype : jnp .dtype = jnp .float32 ,
934+ mask_padding_tokens : bool = True ,
935+ residual_checkpoint_name : str | None = None ,
936+ attention_mask : jax .Array = None ,
937+ ulysses_shards : int = - 1 ,
938+ use_base2_exp : bool = True ,
939+ use_experimental_scheduler : bool = False ,
940+ bidirectional : bool = False ,
941+ use_fixed_m : bool = False ,
942+ ) -> jax .Array :
943+ """Hybrid Ulysses + Ring (USP) with the CUSTOM splash kernel on main's mesh.
944+
945+ Uses origin/main's explicit internal `(ring, ulysses)` mesh
946+ (`_create_internal_ulysses_ring_mesh`, commit c104db51) instead of single-axis
947+ collective sub-groups: the public `context` axis is reshaped with the Ulysses
948+ axis innermost, so the Ulysses all-to-all stays INTRA-chip and the ring rotates
949+ ACROSS chips. The per-shard attention is our custom splash kernel
950+ (`make_custom_ring_attention`), not the tokamax_ring kernel main uses.
951+
952+ 1. all-to-all over the (intra-chip) Ulysses axis: trade sequence for heads;
953+ 2. ring (full ppermute) over the (cross-chip) ring axis, online-softmax merge;
954+ 3. all-to-all back to restore the sequence-sharded / full-heads layout.
955+
956+ U = ulysses_shards (from config); R = context // U. U=context -> pure
957+ Ulysses, U=1 -> pure Ring (all on the same custom kernel).
958+ """
959+ if attention_mask is not None :
960+ raise NotImplementedError (
961+ "ulysses_ring_custom does not support attention_mask (the custom splash kernels only "
962+ "handle padding via orig_seq_len); got a non-None attention_mask."
963+ )
964+ axis_name = "context"
965+ num_context_shards = mesh .shape [axis_name ]
966+ num_ulysses_shards = ulysses_shards
967+ if num_ulysses_shards <= 0 :
968+ raise ValueError ("ulysses_ring_custom requires ulysses_shards to be set from config or command line." )
969+ if num_context_shards % num_ulysses_shards != 0 :
970+ raise ValueError (
971+ f"ulysses_ring_custom requires ulysses_shards to divide the context shard count, "
972+ f"got context_shards={ num_context_shards } { num_ulysses_shards }
973+ )
974+ num_ring_shards = num_context_shards // num_ulysses_shards
975+
976+ query , orig_q_seq_len = _reshape_data_for_flash (query , heads , num_context_shards )
977+ key , _ = _reshape_data_for_flash (key , heads , num_context_shards )
978+ value , _ = _reshape_data_for_flash (value , heads , num_context_shards )
979+ num_heads = query .shape [1 ]
980+ if num_heads % num_ulysses_shards != 0 :
981+ raise ValueError (f"Ulysses+Ring requires heads divisible by U={ num_ulysses_shards } { num_heads } )
982+
983+ bq , bkv , bkv_compute , bkv_compute_in , heads_per_tile , vmem_limit_bytes = _extract_custom_block_sizes (flash_block_sizes )
984+ if heads_per_tile > 1 :
985+ raise NotImplementedError ("ulysses_ring_custom currently supports heads_per_tile == 1 only." )
986+
987+ internal_mesh = _create_internal_ulysses_ring_mesh (mesh , num_ring_shards , num_ulysses_shards )
988+ ring_axis = INTERNAL_RING_AXIS
989+ ulysses_axis = INTERNAL_ULYSSES_AXIS
990+
991+ q_axis_names = nn .logical_to_mesh_axes (axis_names_q )
992+ kv_axis_names = nn .logical_to_mesh_axes (axis_names_kv )
993+ internal_q_axis_names = _replace_mesh_axis_names (q_axis_names , axis_name , (ring_axis , ulysses_axis ))
994+ internal_kv_axis_names = _replace_mesh_axis_names (kv_axis_names , axis_name , (ring_axis , ulysses_axis ))
995+
996+ @functools .partial (
997+ jax .shard_map ,
998+ mesh = internal_mesh ,
999+ in_specs = (internal_q_axis_names , internal_kv_axis_names , internal_kv_axis_names ),
1000+ out_specs = internal_q_axis_names ,
1001+ check_vma = False ,
1002+ )
1003+ def wrap_ulysses_ring_attention (query , key , value ):
1004+ # (1) Ulysses all-to-all over the (intra-chip) ulysses axis: heads -> sequence,
1005+ # so each device holds the full ring-chunk sequence with heads/U heads.
1006+ a2a = functools .partial (jax .lax .all_to_all , axis_name = ulysses_axis , tiled = True )
1007+ query = a2a (query , split_axis = 1 , concat_axis = 2 )
1008+ key = a2a (key , split_axis = 1 , concat_axis = 2 )
1009+ value = a2a (value , split_axis = 1 , concat_axis = 2 )
1010+
1011+ if use_base2_exp :
1012+ query = query * LOG2E
1013+
1014+ if use_fixed_m :
1015+ # K-smoothing precondition for fixed-m, computed PER SHARD (no ring pmean).
1016+ # A global mean would be a perfectly-uniform per-query logit shift, but the
1017+ # per-shard local mean differs from it by only O(1/sqrt(local_seq)), and the
1018+ # ring's outer online-softmax merge re-normalizes across shards anyway, so we
1019+ # drop the per-layer ring collective and accept the negligible shift error.
1020+ kbar = jnp .mean (key , axis = 2 , keepdims = True )
1021+ key = key - kbar
1022+
1023+ query , kv_size , query_seq_len = _pad_data_for_flash (query , heads , bq )
1024+ key , _ , key_seq_len = _pad_data_for_flash (key , heads , bkv )
1025+ value , _ , _ = _pad_data_for_flash (value , heads , bkv )
1026+
1027+ mk_arr = None
1028+ if use_fixed_m :
1029+ # Per-(local-)head Cauchy-Schwarz inputs, all LOCAL to this ring shard. The
1030+ # outer ring merge does an online softmax across shards, so each shard's
1031+ # kernel may use its own local max||k|| as the fixed-m bound for its own
1032+ # local keys -- no global ring pmax is needed for correctness. This removes
1033+ # the second per-layer ring collective.
1034+ qf = query .astype (jnp .float32 )
1035+ kf = key .astype (jnp .float32 )
1036+ qn_max = jnp .sqrt ((qf * qf ).sum (- 1 )).max (axis = (0 , 2 )) # (local_heads,)
1037+ mk_h = jnp .sqrt ((kf * kf ).sum (- 1 )).max (axis = (0 , 2 )) # (local_heads,) local
1038+ fixed_ok = (qn_max * mk_h <= custom_splash ._FIXED_M_SAFE_BOUND ).astype (jnp .float32 )
1039+ if os .environ .get ("FIXED_M_FORCE_ALL" , "0" ) == "1" :
1040+ # PERF PROBE ONLY (unsafe): force every head onto the fixed-m fast path,
1041+ # bypassing the safety gate, to measure fixed-m's speed CEILING on the
1042+ # ring kernel. Output may be garbage; timing is still valid.
1043+ fixed_ok = jnp .ones_like (fixed_ok )
1044+ mk_arr = jnp .stack ([mk_h , fixed_ok ]) # (2, local_heads)
1045+
1046+ bsizes = custom_splash ._BlockSizes (block_q = bq , block_kv = bkv , block_kv_compute = bkv_compute )
1047+ if num_ring_shards == 1 :
1048+ # (2a) R=1: the ring is trivial (no rotation) -> use the lighter dedicated
1049+ # splash kernel (fuse_reciprocal, no fp32 online-softmax residual windows).
1050+ # Same math as the 1-step ring, and it fits BQ=8448 where the ring kernel
1051+ # OOMs (its 3x residual windows). make_splash_mha returns [H, D, S].
1052+ splash_kernel = custom_splash .make_splash_mha (
1053+ block_sizes = bsizes ,
1054+ bkv_compute_in = bkv_compute_in ,
1055+ orig_q_seq_len = query_seq_len ,
1056+ orig_kv_seq_len = key_seq_len ,
1057+ heads_per_tile = heads_per_tile ,
1058+ use_base2_exp = use_base2_exp ,
1059+ use_experimental_scheduler = use_experimental_scheduler ,
1060+ vmem_limit_bytes = vmem_limit_bytes ,
1061+ use_fixed_m = use_fixed_m ,
1062+ )
1063+ if use_fixed_m :
1064+ attention_output = jnp .swapaxes (jax .vmap (splash_kernel , in_axes = (0 , 0 , 0 , None ))(query , key , value , mk_arr ), 2 , 3 )
1065+ else :
1066+ attention_output = jnp .swapaxes (jax .vmap (splash_kernel , in_axes = (0 , 0 , 0 ))(query , key , value ), 2 , 3 )
1067+ else :
1068+ # (2b) Ring (full ppermute over the cross-chip ring axis) with the custom kernel.
1069+ # bidirectional=True -> wrap-free schedule (streams K/V both directions one hop
1070+ # at a time), for a non-wrapping ring axis. Selected by attention=ulysses_ring_custom_bidir.
1071+ ring_kernel = tokamax_ring_attention_kernel .make_custom_ring_attention (
1072+ block_sizes = bsizes ,
1073+ bkv_compute_in = bkv_compute_in ,
1074+ orig_q_seq_len = query_seq_len ,
1075+ orig_kv_seq_len = key_seq_len ,
1076+ use_base2_exp = use_base2_exp ,
1077+ use_experimental_scheduler = use_experimental_scheduler ,
1078+ vmem_limit_bytes = vmem_limit_bytes ,
1079+ ring_axis = ring_axis ,
1080+ ring_size = num_ring_shards ,
1081+ bidirectional = bidirectional ,
1082+ use_fixed_m = use_fixed_m ,
1083+ mk = mk_arr ,
1084+ )
1085+ attention_output = jax .vmap (ring_kernel , in_axes = (0 , 0 , 0 ))(query , key , value )
1086+ attention_output = attention_output [:, :, :query_seq_len , :kv_size ].astype (query .dtype )
1087+
1088+ # (3) Ulysses all-to-all back: sequence -> heads, restoring the layout.
1089+ attention_output = a2a (attention_output , split_axis = 2 , concat_axis = 1 )
1090+ return attention_output
1091+
1092+ x = wrap_ulysses_ring_attention (query , key , value )
1093+ x = jax .lax .with_sharding_constraint (x , q_axis_names )
1094+ x = x [:, :, :orig_q_seq_len , :]
1095+ x = _reshape_heads_to_head_dim (x )
1096+ return x
1097+
1098+
8881099def _apply_attention_dot (
8891100 query : Array ,
8901101 key : Array ,
0 commit comments