Fix ROCm MLA LSE layout; preserve CIRU refs

.gitattributes

@@ -4,4 +4,4 @@
 *.md text eol=lf
 *.png filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-Ling-3.0-Flash-CIRU-int4-Strix-native.tar.gz filter=lfs diff=lfs merge=lfs -text
+runtime/Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.bundle filter=lfs diff=lfs merge=lfs -text

CHANGELOG.md

@@ -6,8 +6,8 @@
 - Model revision: `ca3ea63b0255d212c4fe6020db9e0a51ce136006`
 - Model format: 24 safetensors shards, 77,012,299,464 bytes total
 - vLLM upstream base: `d35eb6c44071ea806018841c490f0d2f3219c485`
-- CIRU runtime head: `388d82de2abe418e73380ba250270d3abc62ac12`
-- Net vLLM delta: six files, 365 insertions, 26 deletions
+- CIRU runtime head: `91fab4e6dcaa683911655630dae8a5b140a7cb91`
+- Net vLLM delta: eight files, 459 insertions, 28 deletions, including regression tests
 
 ## Weight changes
 
@@ -23,6 +23,10 @@ Uses 512 rather than 1024 merge threads on ROCm so Wave32 hardware does not exce
 
 Adds an opt-in Torch implementation selected by `VLLM_ROCM_SAFE_MERGE_ATTN_STATES=1`. It handles token-first and head-first LSE layouts and avoids an observed `gfx1151` HSA fault in the Triton merge kernel.
 
+### `vllm/v1/attention/backends/mla/prefill/flash_attn.py`
+
+Normalizes the LSE returned by upstream ROCm FlashAttention variable-length prefill from token-first `[tokens, heads]` to vLLM's documented head-first `[heads, tokens]` adapter contract. This prevents heterogeneous chunked-prefill batches from copying a head extent into a token extent in the MLA context accumulator.
+
 ### `vllm/v1/attention/backends/mla/triton_mla.py`
 
 Declares uniform query-length support and converts causal multi-token verifier blocks into per-token Triton MLA decode rows with correct causal KV-prefix lengths. This is relevant to MLA speculative verification, not W4A16 specifically.
@@ -39,6 +43,10 @@ Adds a low-launch-overhead Triton reduction for one to three tokens, top-k 8, an
 
 Wires the guarded WNA16 paths together and adds the exact small-shape `gfx1151` SiLU-and-multiply kernel. Every adjacent shape or unsupported configuration retains the upstream route.
 
+### `tests/v1/attention/test_mla_prefill_quant_output.py`
+
+Adds adapter-layout coverage for both upstream ROCm FlashAttention and vLLM FlashAttention, plus an adapter-to-accumulator regression that exercises the formerly failing token/head extent mismatch.
+
 ## Runtime configuration
 
 - Triton MLA attention
@@ -68,3 +76,4 @@ The final exact-shape SiLU candidate passed bitwise source fixtures and exact ou
 - Added an explicitly experimental Windows 11 WSL2/AMD ROCDXG path; native Windows vLLM is not claimed.
 - Initially changed the public native-256K profile from `gpu_memory_utilization=0.82` to the measured `0.75` envelope used by the retained 60K benchmark; the release default was subsequently reduced to `0.72` as recorded above.
 - Added the retained 60K PP/TG comparison against AtomicChat AD-IQ4_XXS and ROCmFP4 STRIX MTP, with the native row marked for a final-commit refresh.
+- Fixed the ROCm MLA chunked-prefill LSE layout mismatch in runtime commit `91fab4e6d`. Focused attention tests passed 31/31 (with one unrelated CUDA-only FA4 test deselected), and the exact five-request tool-calling reproducer completed 5/5 scenarios with zero errors while the original server PID remained alive.

MANIFEST.json

@@ -22,11 +22,11 @@
   "runtime": {
     "upstream_repo": "https://github.com/vllm-project/vllm.git",
     "upstream_base": "d35eb6c44071ea806018841c490f0d2f3219c485",
-    "ciru_head": "388d82de2abe418e73380ba250270d3abc62ac12",
-    "commits_ahead": 15,
-    "net_files_changed": 6,
-    "net_insertions": 365,
-    "net_deletions": 26
+    "ciru_head": "91fab4e6dcaa683911655630dae8a5b140a7cb91",
+    "commits_ahead": 16,
+    "net_files_changed": 8,
+    "net_insertions": 459,
+    "net_deletions": 28
   },
   "validated_stack": {
     "python": "3.12.12",
@@ -60,7 +60,17 @@
       "c4_aggregate_decode_tokens_per_second": 59.03,
       "c5_aggregate_decode_tokens_per_second": 62.16,
       "validated_max_concurrency": 5,
-      "c6_status": "engine failure in chunked-prefill MLA accumulation; public profiles cap active sequences below this boundary"
+      "c6_status": "not rerun after the LSE layout fix; no post-fix C6 claim"
+    },
+    "rocm_mla_lse_fix_validation": {
+      "runtime_commit": "91fab4e6dcaa683911655630dae8a5b140a7cb91",
+      "focused_tests_passed": 31,
+      "cuda_only_tests_deselected": 1,
+      "reproducer_concurrency": 5,
+      "reproducer_scenarios_passed": 5,
+      "reproducer_request_error_rate": 0.0,
+      "server_process_survived": true,
+      "result_sha256": "f819148f8d49342a73fac2c8be6ce6cd08dfca01c238e94dd586bbd6e5894cd7"
     },
     "retained_60k_prompt_512_output": {
       "native_source": "fabbc7e928dfa64c29dcdf0f795b05246555f4ae",

README.md

@@ -28,10 +28,11 @@ The model weights were **not** changed, requantized, merged, fine-tuned, or rena
 
 ## What this release changes
 
-The runtime is based on vLLM commit `d35eb6c44071ea806018841c490f0d2f3219c485` plus a 15-commit CIRU branch ending at `388d82de2abe418e73380ba250270d3abc62ac12`. Its net source delta is six files, 365 insertions, and 26 deletions.
+The runtime is based on vLLM commit `d35eb6c44071ea806018841c490f0d2f3219c485` plus a 16-commit CIRU branch ending at `91fab4e6dcaa683911655630dae8a5b140a7cb91`. Its net source delta is eight files, 459 insertions, and 28 deletions, including regression tests.
 
 - Includes [`liminfei-amd`'s Wave32 LDS fix](https://github.com/vllm-project/vllm/pull/46012) for the ROCm top-k merge launch, preserving the AMD author's commit credit.
 - Adds an opt-in safe PyTorch attention-state merge for a `gfx1151` Triton/HSA fault, including both supported MLA LSE layouts.
+- Normalizes the upstream ROCm FlashAttention variable-length LSE result to vLLM's documented head-first layout before chunked MLA context accumulation.
 - Routes causal uniform multi-token MLA verification blocks through Triton decode with correct per-token sequence lengths.
 - Adds narrowly guarded W4A16 MoE decode/verification fast paths for Ling's exact small-token geometry.
 - Adds a `gfx1151` Triton SiLU-and-multiply kernel for Ling's exact eager W4A16 shapes.
@@ -111,7 +112,7 @@ The named release archive is `Ling-3.0-Flash-CIRU-int4-Strix-native.tar.gz`. It
 
 ## Quick start
 
-The exact validated environment is Python 3.12, Torch `2.13.0+rocm7.15.0a20260728`, HIP `7.15.0`, Triton `3.8.0+git4cff872c.rocm7.15.0a20260728`, FlashAttention `2.8.3`, and the CIRU vLLM source at `388d82de2`.
+The exact validated environment is Python 3.12, Torch `2.13.0+rocm7.15.0a20260728`, HIP `7.15.0`, Triton `3.8.0+git4cff872c.rocm7.15.0a20260728`, FlashAttention `2.8.3`, and the CIRU vLLM source at `91fab4e6d`.
 
 For a guided install, run this from the repository root:
 
@@ -189,7 +190,7 @@ The exact runtime has been validated on CIRU's native-Linux NixOS Strix Halo hos
 
 The official checkpoint is natively configured and trained for 256K context. `run-256k.sh` uses `max_model_len=262144`, `max_num_seqs=5`, and `gpu_memory_utilization=0.72`. On CIRU's 128 GB Strix Halo host, this exact profile reported a 1,552,168-token KV pool and 5.92x capacity at 262,144 tokens per request, enough for five full native-context sequences.
 
-Five is also the validated serving safety cap. A six-request 2K+256 simultaneous burst terminated the engine in the chunked-prefill MLA accumulation path with an LSE extent mismatch. With `max_num_seqs=5`, additional requests wait in the scheduler instead of entering that unvalidated C6 path.
+Five is the packaged concurrency setting and the validated simultaneous serving load. The earlier MLA accumulator failure was reproduced with five heterogeneous tool-calling requests and root-caused to the ROCm upstream FlashAttention LSE result arriving as `[tokens, heads]` while vLLM's chunk accumulator requires `[heads, tokens]`; it was not caused by a request occupying “slot 5.” Runtime commit `91fab4e6d` normalizes that adapter output. The exact five-request reproducer then completed 5/5 scenarios with zero request errors and the same server process remained healthy. Six-way serving has not been rerun after this fix, so this release makes no C6 stability claim. Requests beyond the configured five wait in vLLM's scheduler.
 
 The 1M profile is an experimental 4x YaRN extension, not native 1M training. It uses `max_model_len=1048576`, `max_num_seqs=2`, and `gpu_memory_utilization=0.82`; further requests wait in the vLLM scheduler.
 
@@ -199,6 +200,7 @@ The fork remains a vLLM runtime and safely falls back to upstream paths when its
 
 - The Wave32 top-k fix is broadly relevant to ROCm Wave32 GPUs.
 - Safe attention merging and uniform verifier routing are relevant to compatible ROCm MLA workloads, independent of W4A16.
+- The upstream ROCm FlashAttention LSE normalization is relevant to compatible ROCm MLA chunked-prefill workloads, independent of Ling's W4A16 geometry.
 - The W4A16 fast assignment path currently requires 512 experts, top-k 8, block size 32, and one to three tokens.
 - The fast reduction requires hidden size 2560 and top-k 8.
 - The SiLU path requires exact `(8,1536)->(8,768)` or `(16,1536)->(16,768)` BF16 shapes on `gfx1151`.

SHA256SUMS

@@ -1,22 +1,22 @@
 7b77dda00cc6bd6d0ce15f7c055f9e997f56e1cf1709c61cca5a57283ce3f008  .gitattributes
-411ac21b4849cb3bf65f84a55263550dcd658e0623c2b117233ba1fced1c17ff  assets/ling30int4.png
-65cf5b48f9df164c40fbac98e21c6f9c5ced98660d653630cb58f5ab4b28919f  CHANGELOG.md
-16213dfa6abe90226ab92f9407f3bb92f9ed729e30bb8c37ffd92c7dbc999c73  config/experimental-yarn-1m.env
-c3cdeda6c1bdd47f2b947eac7c4c7182e27249220ddd6eb372c5d12c2c6c405f  config/native-256k.env
-0bb41680a6f0c8366e8f883ca7d76ce1a2ccd3e42541f6d0616df33fe867f55a  docs/BENCHMARKS.md
+411ac21b4849cb3bf65f84a55263550dcd658e0623c2b117233ba1fced1c17ff  assets/ling30int4.png
+45bdb22fa81cdd2d41b111060db97ab6eebdf176d37eba87065ffe6e8d49843c  CHANGELOG.md
+16213dfa6abe90226ab92f9407f3bb92f9ed729e30bb8c37ffd92c7dbc999c73  config/experimental-yarn-1m.env
+c3cdeda6c1bdd47f2b947eac7c4c7182e27249220ddd6eb372c5d12c2c6c405f  config/native-256k.env
+658f33987b057ac0d62a9bad61eb437187d1509e1275424778ddf570476db112  docs/BENCHMARKS.md
 500a1953f2e27379ff90a70db485251c781290a49a3837fd91783e04adadd597  docs/UBUNTU_BUILD.md
-23336099c40774a77a555165da7aee562d5e603017201668178f35f6a8e5cd0d  docs/VALIDATED_ENVIRONMENT.md
-3176eba46cccbd60d33589230190e530eaace8c96466c9621772a1f2c50c07e4  docs/WINDOWS_WSL2.md
+7014d768a0cd85681404e2d571325ae250f6a0713b4ca78c08d69a29ee8916d4  docs/VALIDATED_ENVIRONMENT.md
+3176eba46cccbd60d33589230190e530eaace8c96466c9621772a1f2c50c07e4  docs/WINDOWS_WSL2.md
 e7128cfb97cfac9d144cb999a45b7506c612a8678cfec7f292e00a23ebb8ca8d  install.sh
-852b2422c20355b533884dae00f330aea8e836efecf0b070d59b3e848540c169  MANIFEST.json
-e0fe2d5b9a74a077f1fc4316fb8b7d2d5c57f5cd2aaf4287cd4518761ed2dbaa  README.md
-15e48a8df3b65a73c81cf73c820ba12b4ab9474697d6b3f876c26c0b0cd3d455  runtime/COMMITS.txt
-6b03af3a1685a1969ff96870ea776537f5f67fd3674fc353880bd088bcef9a9c  runtime/Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.bundle
-8f307e9f2f16207856ca0b064a765a73bbebc029ad0b0fed61e5975a17e07ab6  runtime/Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.patch
-f4caeda80ed751a40fa976aa7853357b88075e6eb6d2d49e68b0bc8599754c34  runtime/SHA256SUMS.runtime
-2c066e583d3b24d4c0caa7d04263228dff80e464c61a557736e9c5c5f79b9869  scripts/build-vllm-gfx1151.sh
+57ec70f2fb867e5de28dce1e3ef9d66ac879f78fc502e95da42d3b79c09478a4  MANIFEST.json
+496be3874793c2902e3baf419c0e3b8941a7b331e480f0dc09f9e228bec2d6fe  README.md
+a14fcd3747829be98a73a453057bd451a492d87a236c2affcf57db1cbbf18386  runtime/COMMITS.txt
+9cbcbbef5e113e22b753801ee141bd877b20c84f25f15cc159d0ae508aaf4324  runtime/Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.bundle
+c0c1446ea2c7334ef8b888f3104ec660d2ab41d2cb08bb9d90a3f589dba040e8  runtime/Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.patch
+ecab16a7d189d48ff85942ebefd49f2c206445a4e4a85de440ebde19b6555551  runtime/SHA256SUMS.runtime
+ca51fd447132732880d523db21527fce6f8126925fb175705640037eaecd5f97  scripts/build-vllm-gfx1151.sh
 e6e987937ae4cb544b0aa9b3c90c29acd99c0a9bb8569d42bc95ea7169be245e  scripts/download-model.sh
 8142ef953e2c9618df885756eb897303e6443a4d139a09fe9a882c80d054f466  scripts/install-host-deps.sh
-27d3cfba9f4178f83cbd2d7e44db098ab69e515f76e5c98c40ac219b7b0597ca  scripts/run-1m-yarn-experimental.sh
-296ee94a925db5b97e21afb6ffd4ae3c0cc1c2be67af81e21350f43efaa59ead  scripts/run-256k.sh
+27d3cfba9f4178f83cbd2d7e44db098ab69e515f76e5c98c40ac219b7b0597ca  scripts/run-1m-yarn-experimental.sh
+296ee94a925db5b97e21afb6ffd4ae3c0cc1c2be67af81e21350f43efaa59ead  scripts/run-256k.sh
 718bd922b3bb0a07fb9e2bb0c5ca9a67b799582ac6f614c17af291793e2a69e0  scripts/run-common.sh

docs/BENCHMARKS.md

@@ -18,7 +18,7 @@ The native row used the then-current `gpu_memory_utilization=0.75`, 262,144-cont
 
 ### Update status
 
-The native 60K row was measured at CIRU runtime commit `fabbc7e928dfa64c29dcdf0f795b05246555f4ae`, immediately before the final exact-shape SiLU kernel commit. The released commit `388d82de2abe418e73380ba250270d3abc62ac12` has not yet been rerun on this 60K protocol. The table is therefore a valid retained comparison, but not a claim that 19.92 tok/s is the final commit's maximum 60K speed. Replace the native row only after a matched rerun of all three stacks or a clearly labeled native-only refresh.
+The native 60K row was measured at CIRU runtime commit `fabbc7e928dfa64c29dcdf0f795b05246555f4ae`, before the final exact-shape SiLU kernel and ROCm MLA LSE-layout commits. The released commit `91fab4e6dcaa683911655630dae8a5b140a7cb91` has not yet been rerun on this 60K protocol. The table is therefore a valid retained comparison, but not a claim that 19.92 tok/s is the final commit's maximum 60K speed. Replace the native row only after a matched rerun of all three stacks or a clearly labeled native-only refresh.
 
 ## Short decode and concurrent serving
 

docs/VALIDATED_ENVIRONMENT.md

@@ -19,7 +19,7 @@
 - torchaudio `2.11.0.2+rocm7.15.0a20260728`
 - FlashAttention `2.8.3`, Dao-AILab commit `0e60e39473e8df549a20fb5353760f7a65b30e2d`, Python-only AMD Triton path
 - vLLM upstream base `d35eb6c44071ea806018841c490f0d2f3219c485`
-- CIRU runtime head `388d82de2abe418e73380ba250270d3abc62ac12`
+- CIRU runtime head `91fab4e6dcaa683911655630dae8a5b140a7cb91`
 
 ## Build facts
 
@@ -34,4 +34,6 @@ The Ubuntu/Debian, Fedora, Arch, and Windows/WSL build paths in this package are
 
 The published native-256K launcher defaults to `gpu_memory_utilization=0.72`, `max_model_len=262144`, and `max_num_seqs=5`. The running CIRU profile exposed a 1,552,168-token KV pool and reported 5.92x maximum concurrency at the full native context. Users who do not need that KV-cache capacity can lower the memory fraction to leave more unified memory free, at the cost of full-context concurrency.
 
+The release head includes a ROCm MLA adapter fix that normalizes upstream FlashAttention LSE output to vLLM's head-first accumulator contract. Focused attention coverage passed 31 tests, with one unrelated CUDA-only FA4 test deliberately deselected. The exact heterogeneous five-request workload that previously exposed the mismatch completed all five scenarios with zero request errors after the fix, and the original server process remained healthy. C6 has not been rerun after the fix.
+
 The retained 60K+512 benchmark was measured earlier at `gpu_memory_utilization=0.75`; that historical measurement completed without faults but left only about 2 GiB of available system memory at peak. Its result is not relabeled as a 0.72 measurement.

runtime/COMMITS.txt

@@ -13,3 +13,4 @@ abfff126e0698d98f1872f87f501d36a19382240  Optimize single-token MoE reduction
 928e725a9dbcdbea0bf6c914c8b3fe93f08b0d3f  Optimize K2 verification MoE support kernels
 fabbc7e928dfa64c29dcdf0f795b05246555f4ae  Extend MoE support fast path to K1 verification
 388d82de2abe418e73380ba250270d3abc62ac12  ROCm: optimize exact small WNA16 SiLU activation
+91fab4e6dcaa683911655630dae8a5b140a7cb91  [Bugfix][ROCm][MLA] Normalize upstream FlashAttention LSE layout

runtime/Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.patch

@@ -1,5 +1,5 @@
 diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu
-index 519e213281cd07a9a4cc81f7bc5fcb843277528d..847d8284afc992c3f75e2d65f240afe3d3d0e0e3 100644
+index 519e21328..847d8284a 100644
 --- a/csrc/libtorch_stable/sampler.cu
 +++ b/csrc/libtorch_stable/sampler.cu
 @@ -711,7 +711,12 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
@@ -15,8 +15,129 @@ index 519e213281cd07a9a4cc81f7bc5fcb843277528d..847d8284afc992c3f75e2d65f240afe3
      vllm::topKPerRowDecode<kNumThreadsPerBlockMerge, true, false, true>
          <<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(
              outLogitsAux.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
+diff --git a/tests/v1/attention/test_mla_prefill_quant_output.py b/tests/v1/attention/test_mla_prefill_quant_output.py
+index d7659485a..93df32d38 100644
+--- a/tests/v1/attention/test_mla_prefill_quant_output.py
++++ b/tests/v1/attention/test_mla_prefill_quant_output.py
+@@ -1,8 +1,9 @@
+ # SPDX-License-Identifier: Apache-2.0
+ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+-"""Tests for MLA prefill backend fused-quant-output support.
++"""Tests for MLA prefill adapter contracts and fused-quant-output support.
+ 
+-Covers two things:
++Covers three things:
++  * ROCm upstream FlashAttention LSE layout normalization at the adapter boundary.
+   * `MLAPrefillBackend.supports_quant_output`, the capability gate that decides
+     whether the prefill kernel writes quantized output directly (FA4 native
+     fused FP8, see flash-attention#135) instead of the post-quant path.
+@@ -10,11 +11,16 @@ Covers two things:
+     + standalone static-FP8-quant path it replaces (GPU-only, SM100/SM110).
+ """
+ 
++from types import SimpleNamespace
+ from unittest.mock import patch
+ 
+ import pytest
+ import torch
+ 
++from vllm.model_executor.layers.attention.mla_attention import (
++    accumulate_mla_context_chunk,
++    init_mla_context_partial,
++)
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
+     kFp8Dynamic128Sym,
+     kFp8StaticTensorSym,
+@@ -57,9 +63,87 @@ def _make_fa_backend(version: int | None, is_vllm_fa: bool):
+     backend = object.__new__(FlashAttnPrefillBackend)
+     backend.vllm_flash_attn_version = version
+     backend._is_vllm_fa = is_vllm_fa
++    backend.flash_attn_varlen_func = None
+     return backend
+ 
+ 
++@pytest.mark.parametrize("is_vllm_fa", [False, True])
++def test_flash_attn_normalizes_rocm_varlen_lse_layout(is_vllm_fa):
++    """ROCm upstream FA returns [tokens, heads]; MLA requires [heads, tokens]."""
++    backend = _make_fa_backend(version=None, is_vllm_fa=is_vllm_fa)
++    backend.requires_v_padding = False
++    tokens, heads, dim = 5, 3, 4
++    attn_out = torch.zeros(tokens, heads, dim)
++    token_major_lse = torch.arange(tokens * heads).reshape(tokens, heads)
++
++    with patch(
++        f"{_FA_MODULE}.FA4_MLA_PREFILL_KERNEL",
++        return_value=(attn_out, token_major_lse),
++    ):
++        _, lse = backend._flash_attn_varlen_diff_headdims(
++            q=attn_out,
++            k=attn_out,
++            v=attn_out,
++            return_softmax_lse=True,
++        )
++
++    expected = token_major_lse if is_vllm_fa else token_major_lse.transpose(0, 1)
++    torch.testing.assert_close(lse, expected)
++    if not is_vllm_fa:
++        assert lse.is_contiguous()
++
++
++def test_rocm_lse_normalization_satisfies_context_accumulator_contract():
++    """Token-major ROCm LSE must reach MLA accumulation as head-major.
++
++    This reproduces the release failure at the adapter/accumulator boundary:
++    without normalization, the accumulator tries to copy a head-count-wide
++    dimension into a token-count-wide destination and terminates EngineCore.
++    """
++    backend = _make_fa_backend(version=None, is_vllm_fa=False)
++    backend.requires_v_padding = False
++    tokens, heads, dim = 5, 3, 4
++    attn_output = torch.arange(tokens * heads * dim).reshape(tokens, heads, dim)
++    token_major_lse = torch.arange(tokens * heads).reshape(tokens, heads)
++
++    with patch(
++        f"{_FA_MODULE}.FA4_MLA_PREFILL_KERNEL",
++        return_value=(attn_output, token_major_lse),
++    ):
++        normalized_output, normalized_lse = (
++            backend._flash_attn_varlen_diff_headdims(
++                q=attn_output,
++                k=attn_output,
++                v=attn_output,
++                return_softmax_lse=True,
++            )
++        )
++
++    metadata = SimpleNamespace(empty_token_slices=[])
++    output, output_lse = init_mla_context_partial(
++        metadata,
++        normalized_output,
++        normalized_lse,
++        num_tokens=tokens,
++    )
++    chunk = SimpleNamespace(
++        token_slice=slice(0, tokens),
++        continuation_token_end=tokens,
++        is_continuation=False,
++    )
++    accumulate_mla_context_chunk(
++        chunk,
++        normalized_output,
++        normalized_lse,
++        output,
++        output_lse,
++    )
++
++    assert output_lse.shape == (heads, tokens)
++    torch.testing.assert_close(output, attn_output)
++    torch.testing.assert_close(output_lse, token_major_lse.transpose(0, 1))
++
++
+ @pytest.mark.parametrize(
+     ("version", "is_vllm_fa", "dc_major", "quant_key", "expected"),
+     [
 diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
-index 9d93dfe54ac1a81d5c1f512bc266b9732506e31a..dd96f6da4cc3e3655bbab58f64ceed0882f2e7c0 100644
+index 9d93dfe54..dd96f6da4 100644
 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
 +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
 @@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
@@ -225,7 +346,7 @@ index 9d93dfe54ac1a81d5c1f512bc266b9732506e31a..dd96f6da4cc3e3655bbab58f64ceed08
 +            # separate function is required for MoE + LoRA
 +            self.moe_sum(intermediate_cache3, output)
 diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
-index be4930052a9a1c883679f7e7415853b2f9c9f667..034c963b7cee65d253e9c27517463c2e59b2d1c0 100644
+index be4930052..034c963b7 100644
 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py
 +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
 @@ -97,6 +97,7 @@ def fused_moe_kernel_gptq_awq(
@@ -294,7 +415,7 @@ index be4930052a9a1c883679f7e7415853b2f9c9f667..034c963b7cee65d253e9c27517463c2e
      )
  
 diff --git a/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py b/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py
-index 768f41db854e0d532d5faeb3379eab1ab8ad5ef4..827f19fdd2117821861e89cb29571875d56257c0 100644
+index 768f41db8..827f19fdd 100644
 --- a/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py
 +++ b/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py
 @@ -7,6 +7,55 @@ from vllm.platforms import current_platform
@@ -353,8 +474,27 @@ index 768f41db854e0d532d5faeb3379eab1ab8ad5ef4..827f19fdd2117821861e89cb29571875
  @triton.jit
  def moe_fused_mul_sum_kernel(
      inputs_ptr,
+diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py
+index 97cb57c18..5d85b08c0 100644
+--- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py
++++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py
+@@ -416,6 +416,14 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
+         if isinstance(attn_out, tuple):
+             attn_out, lse = attn_out[0], attn_out[1]
+ 
++        # The upstream FlashAttention Triton backend on ROCm returns varlen
++        # LSE as [total_q, num_heads], despite documenting (and fake-tensor
++        # modeling) [num_heads, total_q]. vLLM's MLA merge contract uses the
++        # latter layout. Normalize at the adapter boundary so multi-request
++        # cached-prefix prefills do not fail while accumulating context chunks.
++        if not self._is_vllm_fa and lse is not None:
++            lse = lse.transpose(0, 1).contiguous()
++
+         # Remain consistent with old `flash_attn_varlen_func` where there
+         # is only one output tensor if `return_softmax_lse` is False.
+         if return_softmax_lse:
 diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py
-index 71a38e95d0ad49fd6a0511cf6613b5cab1f62af0..641f68b166ff4b4269eb7ec6f07b5cead3849a38 100644
+index 71a38e95d..641f68b16 100644
 --- a/vllm/v1/attention/backends/mla/triton_mla.py
 +++ b/vllm/v1/attention/backends/mla/triton_mla.py
 @@ -13,6 +13,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
@@ -435,7 +575,7 @@ index 71a38e95d0ad49fd6a0511cf6613b5cab1f62af0..641f68b166ff4b4269eb7ec6f07b5cea
  
          # Run MQA — always pass layer scales. When KV cache is
 diff --git a/vllm/v1/attention/ops/merge_attn_states.py b/vllm/v1/attention/ops/merge_attn_states.py
-index 20c7503e9d0518a1c9459c6a599708f94a2bae62..cc860e1ab26ca78bf747a00b160314c4d740c5b1 100644
+index 20c7503e9..cc860e1ab 100644
 --- a/vllm/v1/attention/ops/merge_attn_states.py
 +++ b/vllm/v1/attention/ops/merge_attn_states.py
 @@ -1,11 +1,87 @@

runtime/SHA256SUMS.runtime

@@ -1,3 +1,3 @@
-15e48a8df3b65a73c81cf73c820ba12b4ab9474697d6b3f876c26c0b0cd3d455  COMMITS.txt
-6b03af3a1685a1969ff96870ea776537f5f67fd3674fc353880bd088bcef9a9c  Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.bundle
-8f307e9f2f16207856ca0b064a765a73bbebc029ad0b0fed61e5975a17e07ab6  Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.patch
+a14fcd3747829be98a73a453057bd451a492d87a236c2affcf57db1cbbf18386  COMMITS.txt
+9cbcbbef5e113e22b753801ee141bd877b20c84f25f15cc159d0ae508aaf4324  Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.bundle
+c0c1446ea2c7334ef8b888f3104ec660d2ab41d2cb08bb9d90a3f589dba040e8  Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.patch

scripts/build-vllm-gfx1151.sh

@@ -8,7 +8,7 @@ venv="$install_root/.venv"
 source_dir="$install_root/vllm"
 index=https://rocm.nightlies.amd.com/whl-multi-arch/
 base=d35eb6c44071ea806018841c490f0d2f3219c485
-head=388d82de2abe418e73380ba250270d3abc62ac12
+head=91fab4e6dcaa683911655630dae8a5b140a7cb91
 bundle="$package_root/runtime/Ling-3.0-Flash-CIRU-int4-Strix-native-vllm.bundle"
 flash_source="$install_root/flash-attention"
 flash_commit=0e60e39473e8df549a20fb5353760f7a65b30e2d
@@ -67,7 +67,7 @@ fi
 
 git -C "$source_dir" fetch origin "$base"
 git -C "$source_dir" fetch "$bundle" \
-  refs/heads/codex/ling3-gfx1151-rocm715:refs/heads/ciru/ling3-gfx1151-rocm715
+  refs/heads/ciru/ling3-gfx1151-rocm715:refs/heads/ciru/ling3-gfx1151-rocm715
 git -C "$source_dir" checkout ciru/ling3-gfx1151-rocm715
 test "$(git -C "$source_dir" rev-parse HEAD)" = "$head"