Add Llama 3 Megatron model support

src/art/megatron/context_parallel/executor.py

@@ -5,7 +5,7 @@
 import torch
 from torch._dynamo import config as dynamo_config
 import torch.distributed as dist
-from torch.nn.attention.flex_attention import AuxOutput, AuxRequest, BlockMask
+from torch.nn.attention.flex_attention import BlockMask
 import triton
 import triton.language as tl
 
@@ -16,13 +16,15 @@
     get_sparse_compiled_flex_attention,
     normalize_flex_lse,
     normalize_sparse_block_size,
+    prepare_sparse_flex_attention,
     select_sparse_execution_family,
     sparse_compiled_flex_attention,
 )
 
 from .block_mask import build_block_mask_from_context, prepare_block_mask_context
 from .comm import A2AVCommunicator
 from .range_ops import (
+    prepare_range_head_major_kernels,
     range_gather_head_major,
     range_reduce_sum_,
     range_reduce_sum_head_major_,
@@ -688,21 +690,21 @@ def run(
                 triton_num_stages_2_head_dims=self.triton_num_stages_2_head_dims,
             )
         )
-        out, aux = cast(
-            tuple[torch.Tensor, AuxOutput],
-            compiled_flex_attention(
-                q,
-                k,
-                v,
-                block_mask=block_mask,
-                scale=scale,
-                enable_gqa=enable_gqa,
-                return_aux=AuxRequest(lse=True),
-            ),
+        prepare_sparse_flex_attention(
+            q,
+            k,
+            v,
+            block_mask=block_mask,
+            enable_gqa=enable_gqa,
+        )
+        out, lse = compiled_flex_attention(
+            q,
+            k,
+            v,
+            block_mask=block_mask,
+            scale=scale,
+            enable_gqa=enable_gqa,
         )
-        lse = aux.lse
-        if lse is None:
-            raise RuntimeError("Compiled flex attention did not return lse.")
         lse = normalize_flex_lse(lse, backend=backend)
         return out, lse
 
@@ -1021,9 +1023,9 @@ def _run_stage_attention(
         head_major=input_head_major,
     )
     if input_head_major:
-        q_flex = q_stage.unsqueeze(0)
-        k_flex = k_stage.unsqueeze(0)
-        v_flex = v_stage.unsqueeze(0)
+        q_flex = q_stage.unsqueeze(0).contiguous()
+        k_flex = k_stage.unsqueeze(0).contiguous()
+        v_flex = v_stage.unsqueeze(0).contiguous()
     else:
         q_flex = q_stage.permute(1, 0, 2).unsqueeze(0).contiguous()
         k_flex = k_stage.permute(1, 0, 2).unsqueeze(0).contiguous()
@@ -1887,11 +1889,14 @@ def _flatten_qkv(
     value: torch.Tensor,
     state: ArtContextParallelState,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    return (
+    flattened = (
         flatten_valid_sequence_head_major(query, state.rank_plan.local_valid_lengths),
         flatten_valid_sequence_head_major(key, state.rank_plan.local_valid_lengths),
         flatten_valid_sequence_head_major(value, state.rank_plan.local_valid_lengths),
     )
+    for tensor in flattened:
+        prepare_range_head_major_kernels(tensor)
+    return flattened
 
 
 def _run_context_parallel_forward(

src/art/megatron/context_parallel/range_ops.py

@@ -8,6 +8,8 @@
 
 from .types import TokenRange
 
+_WARMED_HEAD_MAJOR_KERNELS: set[tuple[str, int | None, torch.dtype, int]] = set()
+
 
 def _single_range(ranges: Sequence[TokenRange]) -> TokenRange | None:
     compact = [range_ for range_ in ranges if range_.size() > 0]
@@ -90,8 +92,7 @@ def _range_gather_head_major_kernel(
     output_head_stride,
     output_token_stride,
     inner_size: tl.constexpr,
-    n_cols: tl.constexpr,
-    n_col_blocks: tl.constexpr,
+    n_cols,
     elem_per_block: tl.constexpr,
 ):
     out_row = tl.program_id(0)
@@ -129,8 +130,7 @@ def _range_reduce_sum_head_major_kernel(
     output_head_stride,
     output_token_stride,
     inner_size: tl.constexpr,
-    n_cols: tl.constexpr,
-    n_col_blocks: tl.constexpr,
+    n_cols,
     elem_per_block: tl.constexpr,
 ):
     in_row = tl.program_id(0)
@@ -196,7 +196,7 @@ def _range_meta(
         dtype=torch.int64,
     )
     cu_range_sizes = torch.cumsum(range_sizes, dim=0)
-    total_size = int(cu_range_sizes[-1].item())
+    total_size = sum(range_.size() for range_ in compact)
     row_map = torch.repeat_interleave(
         torch.arange(len(compact), device=device, dtype=torch.int64),
         range_sizes[1:],
@@ -382,14 +382,46 @@ def _range_gather_head_major_impl(
         output.stride(0),
         output.stride(1),
         inner_size=inner_size,  # ty: ignore[invalid-argument-type]
-        n_cols=n_cols,  # ty: ignore[invalid-argument-type]
-        n_col_blocks=n_col_blocks,
+        n_cols=n_cols,
         elem_per_block=elem_per_block,  # ty: ignore[invalid-argument-type]
         num_warps=4,  # ty: ignore[unknown-argument]
     )
     return output
 
 
+def prepare_range_head_major_kernels(input_tensor: torch.Tensor) -> None:
+    """Compile one shape-polymorphic gather/reduce pair before variable CP plans.
+
+    Prefix trees may not require multi-range Q/KV movement until a late batch.
+    Warm once per device, dtype, and inner dimension so that workload variation
+    cannot expose compilation in a later optimizer step. This enqueues two tiny
+    kernels on the current stream and must remain free of CUDA synchronization.
+    """
+    if not input_tensor.is_cuda:
+        return
+    if input_tensor.ndim < 3:
+        raise RuntimeError(
+            "Head-major range kernel preparation expects [H, T, ...], "
+            f"got {tuple(input_tensor.shape)}"
+        )
+    inner_size = input_tensor.numel() // max(
+        int(input_tensor.shape[0] * input_tensor.shape[1]), 1
+    )
+    key = (
+        input_tensor.device.type,
+        input_tensor.device.index,
+        input_tensor.dtype,
+        inner_size,
+    )
+    if key in _WARMED_HEAD_MAJOR_KERNELS:
+        return
+    probe = input_tensor.new_empty((input_tensor.shape[0], 2, *input_tensor.shape[2:]))
+    ranges = (TokenRange(start=0, end=1), TokenRange(start=1, end=2))
+    gathered = _range_gather_head_major_impl(probe, ranges)
+    range_reduce_sum_head_major_(gathered, output_tensor=probe, ranges=ranges)
+    _WARMED_HEAD_MAJOR_KERNELS.add(key)
+
+
 class _RangeGatherFn(torch.autograd.Function):
     @staticmethod
     def forward(
@@ -675,8 +707,7 @@ def range_reduce_sum_head_major_(
         output_tensor.stride(0),
         output_tensor.stride(1),
         inner_size=inner_size,  # ty: ignore[invalid-argument-type]
-        n_cols=n_cols,  # ty: ignore[invalid-argument-type]
-        n_col_blocks=n_col_blocks,
+        n_cols=n_cols,
         elem_per_block=elem_per_block,  # ty: ignore[invalid-argument-type]
         num_warps=4,  # ty: ignore[unknown-argument]
     )

src/art/megatron/flex_attn/compiled.py

@@ -4,8 +4,12 @@
 from typing import Any, Literal, TypeAlias, cast
 
 import torch
+from torch._higher_order_ops.flex_attention import (
+    flex_attention as flex_attention_hop,
+)
 from torch.nn.attention.flex_attention import (
     AuxRequest,
+    BlockMask,
     FlexKernelOptions,
     flex_attention,
 )
@@ -105,26 +109,72 @@ def _forced_flex_attention_dense(
     )
 
 
-def _forced_flex_attention_sparse(
-    q,
-    k,
-    v,
+def _identity_score(score, _batch, _head, _query, _key):
+    return score
+
+
+def _sparse_flex_attention_with_options(kernel_options: FlexKernelOptions) -> Any:
+    resolved_options = dict(kernel_options)
+    resolved_options.setdefault("PRESCALE_QK", False)
+    resolved_options.setdefault("ROWS_GUARANTEED_SAFE", False)
+    resolved_options.setdefault("BLOCKS_ARE_CONTIGUOUS", False)
+    resolved_options.setdefault("WRITE_DQ", True)
+    resolved_options["OUTPUT_LOGSUMEXP"] = True
+    resolved_options["OUTPUT_MAX"] = False
+
+    def _flex_attention(q, k, v, *, block_mask, scale, enable_gqa):
+        del enable_gqa
+        for tensor in (q, k, v):
+            torch._dynamo.mark_static(tensor, -3)
+            torch._dynamo.mark_static(tensor, -1)
+        out, lse, _max_scores = flex_attention_hop(
+            q,
+            k,
+            v,
+            _identity_score,
+            (q.shape[-2], k.shape[-2], *block_mask.as_tuple()[2:]),
+            scale,
+            resolved_options,
+        )
+        return out, lse * _FLASH_LSE_RESCALE
+
+    return _flex_attention
+
+
+def prepare_sparse_flex_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
     *,
-    block_mask,
-    scale,
-    enable_gqa,
-    return_aux: AuxRequest | None = None,
-):
-    return flex_attention(
-        q,
-        k,
-        v,
-        block_mask=block_mask,
-        scale=scale,
-        enable_gqa=enable_gqa,
-        kernel_options=_FORCED_FLEX_KERNEL_OPTIONS,
-        return_aux=return_aux,
-    )
+    block_mask: BlockMask,
+    enable_gqa: bool,
+) -> None:
+    """Validate ART's sparse boundary and keep stage lengths symbolic."""
+    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
+        raise RuntimeError("Sparse flex attention requires rank-4 q, k, and v.")
+    if int(q.shape[0]) != int(k.shape[0]) or k.shape[:-1] != v.shape[:-1]:
+        raise RuntimeError("Sparse flex attention received incompatible q, k, and v.")
+    if int(q.shape[-1]) != int(k.shape[-1]):
+        raise RuntimeError("Sparse flex attention requires equal q/k head dimensions.")
+    q_heads, kv_heads = int(q.shape[1]), int(k.shape[1])
+    if (q_heads != kv_heads and not enable_gqa) or q_heads % kv_heads != 0:
+        raise RuntimeError(
+            f"Sparse flex attention received q_heads={q_heads}, kv_heads={kv_heads}, "
+            f"enable_gqa={enable_gqa}."
+        )
+    if tuple(block_mask.shape[-2:]) != (int(q.shape[2]), int(k.shape[2])):
+        raise RuntimeError("Sparse flex attention block-mask lengths do not match q/k.")
+    for tensor in (q, k, v):
+        torch._dynamo.mark_dynamic(tensor, 2)
+    for value in block_mask.as_tuple()[2:]:
+        if isinstance(value, torch.Tensor):
+            for dim in range(2, value.ndim):
+                torch._dynamo.mark_dynamic(value, dim)
+    for cell in getattr(block_mask.mask_mod, "__closure__", None) or ():
+        value = cell.cell_contents
+        if isinstance(value, torch.Tensor):
+            for dim in range(value.ndim):
+                torch._dynamo.mark_dynamic(value, dim)
 
 
 def _flex_attention_with_options(kernel_options: FlexKernelOptions) -> Any:
@@ -161,11 +211,16 @@ def select_sparse_execution_family(
 ) -> tuple[int, int, str]:
     del is_local_stage
     q_block, k_block = normalize_sparse_block_size(block_size)
+    # Avoid Flex's separate zero/one-block Dynamo specialization.
     target_q_len = (
-        0 if int(q_len) <= 0 else ((int(q_len) + q_block - 1) // q_block) * q_block
+        0
+        if int(q_len) <= 0
+        else max(2, (int(q_len) + q_block - 1) // q_block) * q_block
     )
     target_k_len = (
-        0 if int(k_len) <= 0 else ((int(k_len) + k_block - 1) // k_block) * k_block
+        0
+        if int(k_len) <= 0
+        else max(2, (int(k_len) + k_block - 1) // k_block) * k_block
     )
     return int(target_q_len), int(target_k_len), "sparse"
 
@@ -243,14 +298,14 @@ def get_sparse_compiled_flex_attention(
 )
 
 sparse_compiled_flex_attention = torch.compile(
-    _forced_flex_attention_sparse,
+    _sparse_flex_attention_with_options(_FORCED_FLEX_KERNEL_OPTIONS),
 )
 flash_sparse_compiled_flex_attention = torch.compile(
-    _flex_attention_with_options(_FLASH_FLEX_KERNEL_OPTIONS),
+    _sparse_flex_attention_with_options(_FLASH_FLEX_KERNEL_OPTIONS),
 )
 triton_sparse_compiled_flex_attention = torch.compile(
-    _flex_attention_with_options(_TRITON_FLEX_KERNEL_OPTIONS),
+    _sparse_flex_attention_with_options(_TRITON_FLEX_KERNEL_OPTIONS),
 )
 triton_num_stages_2_sparse_compiled_flex_attention = torch.compile(
-    _flex_attention_with_options(_TRITON_NUM_STAGES_2_FLEX_KERNEL_OPTIONS),
+    _sparse_flex_attention_with_options(_TRITON_NUM_STAGES_2_FLEX_KERNEL_OPTIONS),
 )

src/art/megatron/model_support/__init__.py

@@ -4,6 +4,8 @@
     GEMMA4_DENSE_SPEC,
     GEMMA4_MOE_MODELS,
     GEMMA4_MOE_SPEC,
+    LLAMA3_DENSE_MODELS,
+    LLAMA3_DENSE_SPEC,
     PROBE_ONLY_MODEL_SUPPORT_SPECS,
     QWEN3_5_DENSE_MODELS,
     QWEN3_5_DENSE_SPEC,
@@ -66,6 +68,8 @@ def __getattr__(name: str):
     "GEMMA4_MOE_MODELS",
     "GEMMA4_MOE_SPEC",
     "LayerFamilyInstance",
+    "LLAMA3_DENSE_MODELS",
+    "LLAMA3_DENSE_SPEC",
     "ModelSupportHandler",
     "ModelSupportSpec",
     "NativeVllmLoraStatus",

src/art/megatron/model_support/handlers/__init__.py

@@ -16,6 +16,14 @@
         "art.megatron.model_support.handlers.default_dense",
         "DefaultMoeHandler",
     ),
+    "LLAMA3_DENSE_HANDLER": (
+        "art.megatron.model_support.handlers.llama3",
+        "LLAMA3_DENSE_HANDLER",
+    ),
+    "Llama3DenseHandler": (
+        "art.megatron.model_support.handlers.llama3",
+        "Llama3DenseHandler",
+    ),
     "QWEN3_DENSE_HANDLER": (
         "art.megatron.model_support.handlers.qwen3_dense",
         "QWEN3_DENSE_HANDLER",

src/art/megatron/model_support/handlers/llama3.py

@@ -0,0 +1,21 @@
+from typing import Any, Sequence
+
+from art.megatron.model_support.handlers.default_dense import DefaultDenseHandler
+from art.megatron.model_support.handlers.qwen3_common import (
+    install_qwen3_text_preprocess_patch,
+    qwen3_forward_kwargs,
+)
+
+
+class Llama3DenseHandler(DefaultDenseHandler):
+    key = "llama3_dense"
+    native_vllm_lora_status = "validated"
+
+    def install_preprocess_patch(self, model_chunks: Sequence[Any]) -> None:
+        install_qwen3_text_preprocess_patch(model_chunks)
+
+    def get_forward_kwargs(self, model: Any, **kwargs: Any) -> dict[str, Any]:
+        return qwen3_forward_kwargs(model, **kwargs)
+
+
+LLAMA3_DENSE_HANDLER = Llama3DenseHandler()