From 0abcc4faf847972b4366c30c9e7cf896c237db14 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 08:28:05 +0000 Subject: [PATCH 01/17] Add Llama 3 Megatron model support --- src/art/megatron/model_support/__init__.py | 4 +++ .../model_support/handlers/__init__.py | 8 +++++ .../megatron/model_support/handlers/llama3.py | 21 +++++++++++++ src/art/megatron/model_support/registry.py | 31 +++++++++++++++++++ .../model_support/test_provider_support.py | 12 +++++++ .../megatron/model_support/test_workflow.py | 1 + .../megatron/model_support/workflow.py | 1 + 7 files changed, 78 insertions(+) create mode 100644 src/art/megatron/model_support/handlers/llama3.py diff --git a/src/art/megatron/model_support/__init__.py b/src/art/megatron/model_support/__init__.py index 7e1a8b374..6a5735363 100644 --- a/src/art/megatron/model_support/__init__.py +++ b/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", diff --git a/src/art/megatron/model_support/handlers/__init__.py b/src/art/megatron/model_support/handlers/__init__.py index a8da0158d..670468a12 100644 --- a/src/art/megatron/model_support/handlers/__init__.py +++ b/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", diff --git a/src/art/megatron/model_support/handlers/llama3.py b/src/art/megatron/model_support/handlers/llama3.py new file mode 100644 index 000000000..242a61f3a --- /dev/null +++ b/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 = "wip" + + 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() diff --git a/src/art/megatron/model_support/registry.py b/src/art/megatron/model_support/registry.py index 0c9926ab6..f5bdbe871 100644 --- a/src/art/megatron/model_support/registry.py +++ b/src/art/megatron/model_support/registry.py @@ -8,6 +8,7 @@ ) _DEFAULT_DENSE_HANDLER_KEY = "default_dense" +_LLAMA3_DENSE_HANDLER_KEY = "llama3_dense" _QWEN3_DENSE_HANDLER_KEY = "qwen3_dense" _QWEN3_MOE_HANDLER_KEY = "qwen3_moe" _QWEN3_5_DENSE_HANDLER_KEY = "qwen3_5_dense" @@ -81,6 +82,30 @@ native_vllm_lora_status=_DISABLED_NATIVE_VLLM_LORA_STATUS, ) +LLAMA3_DENSE_SPEC = ModelSupportSpec( + key="llama3_dense", + handler_key=_LLAMA3_DENSE_HANDLER_KEY, + model_names=( + "meta-llama/Meta-Llama-3-8B", + "meta-llama/Meta-Llama-3-8B-Instruct", + "meta-llama/Meta-Llama-3-70B", + "meta-llama/Meta-Llama-3-70B-Instruct", + "meta-llama/Llama-3.1-8B", + "meta-llama/Llama-3.1-8B-Instruct", + "meta-llama/Llama-3.1-70B", + "meta-llama/Llama-3.1-70B-Instruct", + "meta-llama/Llama-3.1-405B", + "meta-llama/Llama-3.1-405B-Instruct", + "meta-llama/Llama-3.2-1B", + "meta-llama/Llama-3.2-1B-Instruct", + "meta-llama/Llama-3.2-3B", + "meta-llama/Llama-3.2-3B-Instruct", + "meta-llama/Llama-3.3-70B-Instruct", + ), + default_target_modules=_DENSE_TARGET_MODULES, + native_vllm_lora_status=_WIP_NATIVE_VLLM_LORA_STATUS, +) + QWEN3_MOE_SPEC = ModelSupportSpec( key="qwen3_moe", handler_key=_QWEN3_MOE_HANDLER_KEY, @@ -212,6 +237,7 @@ ) VALIDATED_MODEL_SUPPORT_SPECS = ( + LLAMA3_DENSE_SPEC, QWEN3_MOE_SPEC, QWEN3_DENSE_SPEC, QWEN3_5_MOE_SPEC, @@ -243,6 +269,10 @@ "art.megatron.model_support.handlers.default_dense", "DEFAULT_DENSE_HANDLER", ), + _LLAMA3_DENSE_HANDLER_KEY: ( + "art.megatron.model_support.handlers.llama3", + "LLAMA3_DENSE_HANDLER", + ), _QWEN3_DENSE_HANDLER_KEY: ( "art.megatron.model_support.handlers.qwen3_dense", "QWEN3_DENSE_HANDLER", @@ -302,6 +332,7 @@ _REGISTERED_BRIDGE_KEYS: set[str] = set() QWEN3_DENSE_MODELS = frozenset(QWEN3_DENSE_SPEC.model_names) +LLAMA3_DENSE_MODELS = frozenset(LLAMA3_DENSE_SPEC.model_names) QWEN3_MOE_MODELS = frozenset(QWEN3_MOE_SPEC.model_names) QWEN3_5_DENSE_MODELS = frozenset(QWEN3_5_DENSE_SPEC.model_names) QWEN3_5_MOE_MODELS = frozenset(QWEN3_5_MOE_SPEC.model_names) diff --git a/tests/integration/megatron/model_support/test_provider_support.py b/tests/integration/megatron/model_support/test_provider_support.py index 708d9f6c9..e8ccf9a07 100644 --- a/tests/integration/megatron/model_support/test_provider_support.py +++ b/tests/integration/megatron/model_support/test_provider_support.py @@ -136,7 +136,19 @@ def test_openpipe_qwen3_14b_instruct_uses_qwen3_dense_support() -> None: assert handler.key == "qwen3_dense" +def test_meta_llama_32_1b_instruct_uses_llama3_dense_support() -> None: + model = "meta-llama/Llama-3.2-1B-Instruct" + spec = get_model_support_spec(model) + handler = get_model_support_handler(model) + + assert spec.key == "llama3_dense" + assert spec.is_moe is False + assert spec.native_vllm_lora_status == "wip" + assert handler.key == "llama3_dense" + + def test_model_support_specs_own_moe_metadata() -> None: + assert model_uses_expert_parallel("meta-llama/Llama-3.2-1B-Instruct") is False assert model_uses_expert_parallel("OpenPipe/Qwen3-14B-Instruct") is False assert model_uses_expert_parallel("Qwen/Qwen3-30B-A3B-Instruct-2507") is True assert model_uses_expert_parallel("Qwen/Qwen3.5-35B-A3B") is True diff --git a/tests/integration/megatron/model_support/test_workflow.py b/tests/integration/megatron/model_support/test_workflow.py index 84b796937..08a01ab02 100644 --- a/tests/integration/megatron/model_support/test_workflow.py +++ b/tests/integration/megatron/model_support/test_workflow.py @@ -73,6 +73,7 @@ def test_build_validation_stage_names_has_fixed_order() -> None: def test_validated_architecture_representative_models_are_fixed() -> None: assert validated_architecture_representative_models() == [ + "meta-llama/Llama-3.2-1B-Instruct", "Qwen/Qwen3-30B-A3B", "Qwen/Qwen3-32B", "Qwen/Qwen3.5-35B-A3B", diff --git a/tests/integration/megatron/model_support/workflow.py b/tests/integration/megatron/model_support/workflow.py index f42d3ed21..3a4350930 100644 --- a/tests/integration/megatron/model_support/workflow.py +++ b/tests/integration/megatron/model_support/workflow.py @@ -62,6 +62,7 @@ ) ALL_VALIDATION_STAGES = (*MANDATORY_VALIDATION_STAGES, *OPTIONAL_VALIDATION_STAGES) ARCHITECTURE_REPRESENTATIVE_MODELS = { + "llama3_dense": "meta-llama/Llama-3.2-1B-Instruct", "qwen3_moe": "Qwen/Qwen3-30B-A3B", "qwen3_dense": "Qwen/Qwen3-32B", "qwen3_5_moe": "Qwen/Qwen3.5-35B-A3B", From 67a388649d23a8ca1d0e73323ff1bba50e69fe77 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 08:29:26 +0000 Subject: [PATCH 02/17] Update architecture workflow ordering --- tests/integration/megatron/model_support/test_workflow.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/integration/megatron/model_support/test_workflow.py b/tests/integration/megatron/model_support/test_workflow.py index 08a01ab02..4b68579cd 100644 --- a/tests/integration/megatron/model_support/test_workflow.py +++ b/tests/integration/megatron/model_support/test_workflow.py @@ -244,7 +244,11 @@ def _build_validation_report( stop_on_failure=True, ) - assert calls == ["Qwen/Qwen3-30B-A3B", "Qwen/Qwen3-32B"] + assert calls == [ + "meta-llama/Llama-3.2-1B-Instruct", + "Qwen/Qwen3-30B-A3B", + "Qwen/Qwen3-32B", + ] assert report.passed is False assert [item.base_model for item in report.reports] == calls From 5423cba388ed662339c506c72835f1bd50df4ca4 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 08:35:46 +0000 Subject: [PATCH 03/17] Restore dense FC2 LoRA export --- src/art/megatron/weights/adapter_export.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/art/megatron/weights/adapter_export.py b/src/art/megatron/weights/adapter_export.py index 664e30a02..bb96b51c7 100644 --- a/src/art/megatron/weights/adapter_export.py +++ b/src/art/megatron/weights/adapter_export.py @@ -15,6 +15,7 @@ SelfAttentionLinearProjLoRA, SelfAttentionLinearQKVLoRA, SharedExpertsLinearFC1LoRA, + SharedExpertsLinearFC2LoRA, ) from art.megatron.weights.param_name_canonicalization import canonical_art_param_name @@ -425,10 +426,10 @@ def add_split_mlp_adapter_weights( ) linear_fc2 = getattr(mlp, "linear_fc2", None) - if isinstance(linear_fc2, SelfAttentionLinearProjLoRA): + if isinstance(linear_fc2, SharedExpertsLinearFC2LoRA): fc2_prefix = f"{base_prefix}.linear_fc2" _set_adapter_weights( adapter_weights_by_base, fc2_prefix, - _simple_adapter_weight(fc2_prefix, linear_fc2.lora), + _simple_adapter_weight(fc2_prefix, linear_fc2.row_parallel_lora.lora), ) From e0a4ad0d8c807f8a15f622d628866770e6e48c34 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 09:31:17 +0000 Subject: [PATCH 04/17] Generalize packing rotary validation --- .../model_support/packing_invariance.py | 33 +++++-------------- 1 file changed, 8 insertions(+), 25 deletions(-) diff --git a/tests/integration/megatron/model_support/packing_invariance.py b/tests/integration/megatron/model_support/packing_invariance.py index ca72a2684..68be89ada 100644 --- a/tests/integration/megatron/model_support/packing_invariance.py +++ b/tests/integration/megatron/model_support/packing_invariance.py @@ -49,19 +49,6 @@ PACKING_INVARIANCE_REPORT_FILENAME = "report.json" PACKING_INVARIANCE_ARTIFACT_SUITE_NAME = "Megatron packing-invariance artifacts" REPO_ROOT = Path(__file__).resolve().parents[4] -_SINGLE_ROTARY_OUTPUT_HANDLER_KEYS = frozenset( - { - "default_dense", - "default_moe", - "qwen3_dense", - "qwen3_moe", - "qwen3_5_dense", - "qwen3_5_moe", - "dsv4", - "gpt_oss_moe", - } -) -_TUPLE_ROTARY_OUTPUT_HANDLER_KEYS = frozenset({"gemma4_dense", "gemma4_moe"}) def _slugify(value: str) -> str: @@ -283,21 +270,18 @@ def _rotary_grouping_check( def _rotary_outputs_for_validation( *, - handler_key: str, preprocess_output: Any, ) -> tuple[torch.Tensor | None, ...]: rotary_output = preprocess_output[1] - if handler_key in _SINGLE_ROTARY_OUTPUT_HANDLER_KEYS: - return ( - cast(torch.Tensor | None, rotary_output) - if torch.is_tensor(rotary_output) - else None, - ) - if handler_key in _TUPLE_ROTARY_OUTPUT_HANDLER_KEYS: - local_rotary, global_rotary = rotary_output - return local_rotary, global_rotary + if rotary_output is None or torch.is_tensor(rotary_output): + return (cast(torch.Tensor | None, rotary_output),) + if isinstance(rotary_output, tuple) and all( + item is None or torch.is_tensor(item) for item in rotary_output + ): + return cast(tuple[torch.Tensor | None, ...], rotary_output) raise RuntimeError( - f"Packed position validation has no rotary output mapping for {handler_key!r}" + "Packed position validation received unsupported rotary outputs: " + f"{type(rotary_output).__name__}" ) @@ -703,7 +687,6 @@ def _run_packing_invariance_worker( row_respected = True row_repeated_count = 0 rotary_outputs = _rotary_outputs_for_validation( - handler_key=runtime.model_support_handler.key, preprocess_output=hooked_output, ) for rotary_output in rotary_outputs: From 6657783038d70a5d9cf416490a55598e37dcbff6 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 10:04:56 +0000 Subject: [PATCH 05/17] Validate Llama native LoRA serving --- src/art/megatron/model_support/handlers/llama3.py | 2 +- src/art/megatron/model_support/registry.py | 2 +- .../integration/megatron/model_support/test_provider_support.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/art/megatron/model_support/handlers/llama3.py b/src/art/megatron/model_support/handlers/llama3.py index 242a61f3a..6ab3783fc 100644 --- a/src/art/megatron/model_support/handlers/llama3.py +++ b/src/art/megatron/model_support/handlers/llama3.py @@ -9,7 +9,7 @@ class Llama3DenseHandler(DefaultDenseHandler): key = "llama3_dense" - native_vllm_lora_status = "wip" + native_vllm_lora_status = "validated" def install_preprocess_patch(self, model_chunks: Sequence[Any]) -> None: install_qwen3_text_preprocess_patch(model_chunks) diff --git a/src/art/megatron/model_support/registry.py b/src/art/megatron/model_support/registry.py index f5bdbe871..5511205a9 100644 --- a/src/art/megatron/model_support/registry.py +++ b/src/art/megatron/model_support/registry.py @@ -103,7 +103,7 @@ "meta-llama/Llama-3.3-70B-Instruct", ), default_target_modules=_DENSE_TARGET_MODULES, - native_vllm_lora_status=_WIP_NATIVE_VLLM_LORA_STATUS, + native_vllm_lora_status=_VALIDATED_NATIVE_VLLM_LORA_STATUS, ) QWEN3_MOE_SPEC = ModelSupportSpec( diff --git a/tests/integration/megatron/model_support/test_provider_support.py b/tests/integration/megatron/model_support/test_provider_support.py index e8ccf9a07..43a641452 100644 --- a/tests/integration/megatron/model_support/test_provider_support.py +++ b/tests/integration/megatron/model_support/test_provider_support.py @@ -143,7 +143,7 @@ def test_meta_llama_32_1b_instruct_uses_llama3_dense_support() -> None: assert spec.key == "llama3_dense" assert spec.is_moe is False - assert spec.native_vllm_lora_status == "wip" + assert spec.native_vllm_lora_status == "validated" assert handler.key == "llama3_dense" From 710fcda2bd5563cee02f8915d274dfd971ffc633 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 10:45:37 +0000 Subject: [PATCH 06/17] Keep identity LoRA creation CUDA-free --- src/art/megatron/service.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/art/megatron/service.py b/src/art/megatron/service.py index fe04050f9..92f93c2aa 100644 --- a/src/art/megatron/service.py +++ b/src/art/megatron/service.py @@ -199,9 +199,6 @@ def _skip_meta_to( adapter_config=final_config, ) del peft_model, model - if torch.cuda.is_available(): - torch.cuda.synchronize() - torch.cuda.empty_cache() @dataclass From 012b7a1e6736ecc5c3e3f683cbd72c0c7466181b Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 11:07:05 +0000 Subject: [PATCH 07/17] Prewarm polymorphic CP range kernels --- src/art/megatron/context_parallel/executor.py | 6 ++- .../megatron/context_parallel/range_ops.py | 47 +++++++++++++++---- 2 files changed, 44 insertions(+), 9 deletions(-) diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index cc1de2a0b..68a167ac6 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -23,6 +23,7 @@ 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_, @@ -1887,11 +1888,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( diff --git a/src/art/megatron/context_parallel/range_ops.py b/src/art/megatron/context_parallel/range_ops.py index a645e9f80..bdfb4e53d 100644 --- a/src/art/megatron/context_parallel/range_ops.py +++ b/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) @@ -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] ) From 678162d618fae7c6fd077ce4e2d0edcb18f1f97d Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 11:40:49 +0000 Subject: [PATCH 08/17] Stabilize sparse Flex stage shapes --- src/art/megatron/flex_attn/compiled.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/art/megatron/flex_attn/compiled.py b/src/art/megatron/flex_attn/compiled.py index b6bd8dfe2..e6b1ca0f1 100644 --- a/src/art/megatron/flex_attn/compiled.py +++ b/src/art/megatron/flex_attn/compiled.py @@ -161,11 +161,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" From 3f0256996fde19f1e9733f2183f11049dfdeea12 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 11:58:19 +0000 Subject: [PATCH 09/17] Normalize sparse Flex input layout --- src/art/megatron/context_parallel/executor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index 68a167ac6..baa90385a 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -1022,9 +1022,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() From cf85c078b2a041a7e59b6fd6b4d009637495eef3 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 12:20:11 +0000 Subject: [PATCH 10/17] Keep sparse Flex stage shapes symbolic --- src/art/megatron/context_parallel/executor.py | 31 +++--- src/art/megatron/flex_attn/compiled.py | 96 ++++++++++++++----- .../megatron/model_support/oracle_worker.py | 26 ++++- 3 files changed, 112 insertions(+), 41 deletions(-) diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index baa90385a..b595de990 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/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,6 +16,7 @@ get_sparse_compiled_flex_attention, normalize_flex_lse, normalize_sparse_block_size, + prepare_sparse_flex_attention, select_sparse_execution_family, sparse_compiled_flex_attention, ) @@ -689,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 diff --git a/src/art/megatron/flex_attn/compiled.py b/src/art/megatron/flex_attn/compiled.py index e6b1ca0f1..00b220d96 100644 --- a/src/art/megatron/flex_attn/compiled.py +++ b/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: @@ -248,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), ) diff --git a/tests/integration/megatron/model_support/oracle_worker.py b/tests/integration/megatron/model_support/oracle_worker.py index 7904c1af5..014871667 100644 --- a/tests/integration/megatron/model_support/oracle_worker.py +++ b/tests/integration/megatron/model_support/oracle_worker.py @@ -505,6 +505,24 @@ def _configure_cuda_precision(case_config: OracleCaseConfig) -> None: torch.set_float32_matmul_precision("highest") +def _sparse_flex_test_call(attention_call: Callable[..., Any]) -> Callable[..., Any]: + from torch.nn.attention.flex_attention import AuxRequest + + def sparse_call(q, k, v, *, block_mask, scale, enable_gqa): + out, aux = attention_call( + q, + k, + v, + block_mask=block_mask, + scale=scale, + enable_gqa=enable_gqa, + return_aux=AuxRequest(lse=True), + ) + return out, aux.lse + + return sparse_call + + @contextmanager def _apply_requested_flex_backend_patch(flex_backend: str | None): if flex_backend is None: @@ -539,7 +557,9 @@ def _apply_requested_flex_backend_patch(flex_backend: str | None): compiled_flex_attention._forced_flex_attention_dense ) compiled_flex_attention.sparse_compiled_flex_attention = torch.compile( - compiled_flex_attention._forced_flex_attention_sparse + compiled_flex_attention._sparse_flex_attention_with_options( + patched_kernel_options + ) ) try: yield @@ -594,7 +614,7 @@ def _fp32_inner_call( _fp32_inner_call ) compiled_flex_attention.sparse_compiled_flex_attention = torch.compile( - _fp32_inner_call + _sparse_flex_test_call(_fp32_inner_call) ) try: yield @@ -678,7 +698,7 @@ def _attention_forward_fp32(self, hidden_states, *args, **kwargs): _fp32_inner_call ) compiled_flex_attention.sparse_compiled_flex_attention = torch.compile( - _fp32_inner_call + _sparse_flex_test_call(_fp32_inner_call) ) setattr(ColumnParallelLinear, "_forward_impl", _column_forward_impl_fp32) setattr(RowParallelLinear, "_forward_impl", _row_forward_impl_fp32) From 7942c053cc8e5847cf6694e73a741234e8a005f4 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 13:19:51 +0000 Subject: [PATCH 11/17] Honor Megatron trainability config contract --- .../megatron/trainability/yes_no_trainability.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/integration/megatron/trainability/yes_no_trainability.py b/tests/integration/megatron/trainability/yes_no_trainability.py index 7343b813a..a3cbd6ae6 100644 --- a/tests/integration/megatron/trainability/yes_no_trainability.py +++ b/tests/integration/megatron/trainability/yes_no_trainability.py @@ -50,7 +50,7 @@ _RESOURCE_STAGE_NAME = Literal["yes_no_trainability", "length_trainability"] -class _TrainKwargs(TypedDict): +class _TrainKwargs(TypedDict, total=False): packed_sequence_length: int @@ -550,6 +550,8 @@ def _variant_packed_sequence_length(variant: _TrainabilityVariant) -> int: def _variant_train_kwargs(variant: _TrainabilityVariant) -> _TrainKwargs: + if variant.backend_name == "megatron": + return {} return {"packed_sequence_length": _variant_packed_sequence_length(variant)} @@ -1035,7 +1037,7 @@ async def run_yes_no_trainability_async( 1e-4, ), loss_fn="cispo", - packed_sequence_length=train_kwargs["packed_sequence_length"], + **train_kwargs, ) await model.log( train_groups, From b92ea77b5d99d6f9165a0b1f2b80843a657d901a Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 13:25:10 +0000 Subject: [PATCH 12/17] Use canonical trainability grad metric --- tests/integration/megatron/trainability/test_config.py | 2 +- .../integration/megatron/trainability/yes_no_trainability.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/integration/megatron/trainability/test_config.py b/tests/integration/megatron/trainability/test_config.py index 5262bb8e7..7b01e71f6 100644 --- a/tests/integration/megatron/trainability/test_config.py +++ b/tests/integration/megatron/trainability/test_config.py @@ -217,7 +217,7 @@ def test_yes_no_trainability_passes_initially_saturated_stable_report() -> None: step=1, eval_reward=0.9375, train_reward=0.875, - train_metrics={"grad_norm": 54.0}, + train_metrics={"loss/grad_norm": 54.0}, ) ], ) diff --git a/tests/integration/megatron/trainability/yes_no_trainability.py b/tests/integration/megatron/trainability/yes_no_trainability.py index a3cbd6ae6..a0718bf27 100644 --- a/tests/integration/megatron/trainability/yes_no_trainability.py +++ b/tests/integration/megatron/trainability/yes_no_trainability.py @@ -1121,7 +1121,9 @@ def yes_no_trainability_passed(report: YesNoTrainabilityReport) -> bool: and report.final_eval_reward is not None and report.final_eval_reward >= report.reward_threshold and bool(report.steps) - and any(step.train_metrics.get("grad_norm", 0.0) > 0.0 for step in report.steps) + and any( + step.train_metrics.get("loss/grad_norm", 0.0) > 0.0 for step in report.steps + ) ) return learned_from_below_threshold or already_saturated_and_stable From 734bb52d13cbd8f0ec207392f7ec9fa0ee92320a Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Tue, 14 Jul 2026 13:26:31 +0000 Subject: [PATCH 13/17] Align trainability config test with Megatron --- tests/integration/megatron/trainability/test_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/megatron/trainability/test_config.py b/tests/integration/megatron/trainability/test_config.py index 7b01e71f6..c68b7e210 100644 --- a/tests/integration/megatron/trainability/test_config.py +++ b/tests/integration/megatron/trainability/test_config.py @@ -129,7 +129,7 @@ def test_megatron_variants_keep_short_packed_sequence_default(monkeypatch) -> No ) assert _variant_packed_sequence_length(variant) == 1024 - assert _variant_train_kwargs(variant) == {"packed_sequence_length": 1024} + assert _variant_train_kwargs(variant) == {} config = _build_internal_config( variant, base_model="Qwen/Qwen3-30B-A3B-Instruct-2507" ) From 81c51af74e4ef7dc008828e59779464162795540 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Thu, 16 Jul 2026 22:51:58 +0000 Subject: [PATCH 14/17] Fix protocol trajectory parity indexing --- .../megatron/train_inf_mismatch/real_path.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/tests/integration/megatron/train_inf_mismatch/real_path.py b/tests/integration/megatron/train_inf_mismatch/real_path.py index 11df05a1a..2c6f1fffd 100644 --- a/tests/integration/megatron/train_inf_mismatch/real_path.py +++ b/tests/integration/megatron/train_inf_mismatch/real_path.py @@ -5,6 +5,7 @@ from contextlib import asynccontextmanager, contextmanager import hashlib import inspect +from itertools import chain import json import os from pathlib import Path @@ -507,9 +508,18 @@ def _choice_score_index( indexed: dict[tuple[int, ...], Choice] = {} for group in trajectory_groups: for trajectory in group: - for item in trajectory.messages_and_choices: - if not isinstance(item, Choice): - continue + for item in chain( + ( + item + for item in trajectory.messages_and_choices + if isinstance(item, Choice) + ), + ( + choice + for exchange in trajectory.exchanges.chat_completions + for choice in exchange.response.choices + ), + ): metadata = choice_moe_routing_metadata(item) if metadata is None: if require_routing_metadata: From 1b28bfa14c9cf8270ac961b37b754351bd101d2f Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Mon, 27 Jul 2026 17:37:39 +0000 Subject: [PATCH 15/17] Build workflow parity trajectories explicitly --- .../megatron/train_inf_mismatch/real_path.py | 57 ++++++++----------- .../test_output_parity_invariants.py | 46 +++++++++++++++ 2 files changed, 69 insertions(+), 34 deletions(-) diff --git a/tests/integration/megatron/train_inf_mismatch/real_path.py b/tests/integration/megatron/train_inf_mismatch/real_path.py index 2c6f1fffd..b97e92fc4 100644 --- a/tests/integration/megatron/train_inf_mismatch/real_path.py +++ b/tests/integration/megatron/train_inf_mismatch/real_path.py @@ -5,7 +5,6 @@ from contextlib import asynccontextmanager, contextmanager import hashlib import inspect -from itertools import chain import json import os from pathlib import Path @@ -402,29 +401,28 @@ async def _rollout( import art messages = [{"role": "user", "content": prompt}] - - async def _request() -> None: - request_kwargs: dict[str, Any] = {} - if extra_body is not None: - request_kwargs["extra_body"] = extra_body - response = await model.openai_client().chat.completions.create( - model=model.get_inference_name(), - messages=messages, - max_tokens=max_completion_tokens, - temperature=0.8, - logprobs=True, - top_logprobs=TOP_K, - **request_kwargs, - ) - if trajectory := art.auto_trajectory(): # ty: ignore[deprecated] - logprobs = response.choices[0].logprobs - trajectory.reward = reward - trajectory.metrics["completion_tokens"] = ( + request_kwargs: dict[str, Any] = {} + if extra_body is not None: + request_kwargs["extra_body"] = extra_body + response = await model.openai_client().chat.completions.create( + model=model.get_inference_name(), + messages=messages, + max_tokens=max_completion_tokens, + temperature=0.8, + logprobs=True, + top_logprobs=TOP_K, + **request_kwargs, + ) + choice = response.choices[0] + logprobs = choice.logprobs + return art.Trajectory( + messages_and_choices=[*messages, choice], + reward=reward, + metrics={ + "completion_tokens": ( len(logprobs.content or []) if logprobs is not None else 0 ) - - return await art.capture_auto_trajectory( # ty: ignore[deprecated] - _request() + }, ) @@ -508,18 +506,9 @@ def _choice_score_index( indexed: dict[tuple[int, ...], Choice] = {} for group in trajectory_groups: for trajectory in group: - for item in chain( - ( - item - for item in trajectory.messages_and_choices - if isinstance(item, Choice) - ), - ( - choice - for exchange in trajectory.exchanges.chat_completions - for choice in exchange.response.choices - ), - ): + for item in trajectory.messages_and_choices: + if not isinstance(item, Choice): + continue metadata = choice_moe_routing_metadata(item) if metadata is None: if require_routing_metadata: diff --git a/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py b/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py index 65c2b449e..1e68277fb 100644 --- a/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py +++ b/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py @@ -1,7 +1,11 @@ from __future__ import annotations import math +from types import SimpleNamespace +from unittest.mock import AsyncMock +from openai.types.chat.chat_completion import Choice +from openai.types.chat.chat_completion_message import ChatCompletionMessage import pytest torch = pytest.importorskip("torch") @@ -28,6 +32,7 @@ _delete_adapter_safetensors_on_pass, _real_path_rollout_mode, _real_path_rollout_weights_mode, + _rollout, ) @@ -192,6 +197,47 @@ def test_real_path_default_generates_16_tokens_per_rollout() -> None: assert RealPathConfig().max_completion_tokens == 16 +@pytest.mark.asyncio +async def test_real_path_rollout_builds_explicit_legacy_trajectory() -> None: + choice = Choice( + index=0, + finish_reason="stop", + message=ChatCompletionMessage(role="assistant", content="answer"), + ) + create = AsyncMock(return_value=SimpleNamespace(choices=[choice])) + model = SimpleNamespace( + get_inference_name=lambda: "test-model", + openai_client=lambda: SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=create)) + ), + ) + + trajectory = await _rollout( + model=model, + prompt="question", + max_completion_tokens=7, + reward=0.75, + extra_body={"fixture": True}, + ) + + assert trajectory.messages_and_choices == [ + {"role": "user", "content": "question"}, + choice, + ] + assert trajectory.exchanges.chat_completions == [] + assert trajectory.reward == 0.75 + assert trajectory.metrics["completion_tokens"] == 0 + create.assert_awaited_once_with( + model="test-model", + messages=[{"role": "user", "content": "question"}], + max_tokens=7, + temperature=0.8, + logprobs=True, + top_logprobs=TOP_K, + extra_body={"fixture": True}, + ) + + def test_real_path_rollout_mode_follows_config() -> None: native_config = TrainInfOutputParityConfig( base_model="Qwen/Qwen3.5-35B-A3B", From 55fa89b4a88e469524e6ca88a7a725e213926420 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Wed, 29 Jul 2026 05:07:12 +0000 Subject: [PATCH 16/17] Remove CP synchronization and layout copies --- src/art/megatron/context_parallel/executor.py | 6 +++--- src/art/megatron/context_parallel/range_ops.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index b595de990..e6aa7964c 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -1023,9 +1023,9 @@ def _run_stage_attention( head_major=input_head_major, ) if input_head_major: - q_flex = q_stage.unsqueeze(0).contiguous() - k_flex = k_stage.unsqueeze(0).contiguous() - v_flex = v_stage.unsqueeze(0).contiguous() + q_flex = q_stage.unsqueeze(0) + k_flex = k_stage.unsqueeze(0) + v_flex = v_stage.unsqueeze(0) else: q_flex = q_stage.permute(1, 0, 2).unsqueeze(0).contiguous() k_flex = k_stage.permute(1, 0, 2).unsqueeze(0).contiguous() diff --git a/src/art/megatron/context_parallel/range_ops.py b/src/art/megatron/context_parallel/range_ops.py index bdfb4e53d..7a0ef8b6d 100644 --- a/src/art/megatron/context_parallel/range_ops.py +++ b/src/art/megatron/context_parallel/range_ops.py @@ -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:], From f8b508107fede2198f82ae2f8b4cc27e26b40e95 Mon Sep 17 00:00:00 2001 From: FurtherAI Date: Wed, 29 Jul 2026 05:34:23 +0000 Subject: [PATCH 17/17] Restore stable sparse Flex input layout --- src/art/megatron/context_parallel/executor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index e6aa7964c..b595de990 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -1023,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()