From 00e57d74f07b02a4a83a5b53938a8e5d075c5cde Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Wed, 12 Aug 2026 17:17:26 +0800 Subject: [PATCH] Bump FlashInfer to 0.6.17 and remove Kimi K3 workarounds (#33997) --- docker/Dockerfile | 52 +- .../kimi_k3/flashinfer-perkz-dcp-0.6.15.txt | 5639 ----------------- docker/kimi_k3/kimi_k3_cu12.Dockerfile | 57 +- docker/kimi_k3/kimi_k3_cu13.Dockerfile | 57 +- .../autoregressive/Moonshotai/Kimi-K3.mdx | 13 +- docs/src/snippets/_playground.jsx | 2 +- .../snippets/configs/moonshotai/kimi-k3.jsx | 10 +- python/pyproject.toml | 2 +- .../jit/csrc/kimi_k3/comm/ar_fusion.cuh | 2 +- .../sglang/kernels/ops/moe/trtllm_gen_moe.py | 528 -- python/sglang/srt/arg_groups/overrides.py | 56 +- python/sglang/srt/entrypoints/engine.py | 2 +- python/sglang/srt/environ.py | 9 - .../layers/attention/flashinfer_backend.py | 1 + .../sglang/srt/layers/quantization/mxfp4.py | 58 +- python/sglang/srt/utils/common.py | 2 +- scripts/ci/cuda/ci_install_kimi_k3.sh | 85 +- test/registered/sampling/test_penalty.py | 3 + .../mem_cache/test_unified_mamba_views.py | 2 +- 19 files changed, 84 insertions(+), 6496 deletions(-) delete mode 100644 docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt delete mode 100644 python/sglang/kernels/ops/moe/trtllm_gen_moe.py diff --git a/docker/Dockerfile b/docker/Dockerfile index dd903aeb0..29ae95de1 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -13,10 +13,7 @@ ARG PIP_DEFAULT_INDEX ARG UBUNTU_MIRROR ARG GITHUB_ARTIFACTORY=github.com ARG INSTALL_FLASHINFER_JIT_CACHE=0 -ARG FLASHINFER_VERSION=0.6.15.post1 -ARG TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip" -ARG TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192" -ARG TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1" +ARG FLASHINFER_VERSION=0.6.17 ARG MOONCAKE_VERSION=0.3.12.post1 ARG MSCCLPP_VERSION=sglang-v0.9.1 #if need other arg please add in MOONCAKE_COMPILE_ARG @@ -25,8 +22,7 @@ ARG MOONCAKE_COMPILE_ARG="-DUSE_HTTP=ON -DUSE_MNNVL=ON -DUSE_CUDA=ON -DWITH_EP=O ENV DEBIAN_FRONTEND=noninteractive \ CUDA_HOME=/usr/local/cuda \ GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ \ - FLASHINFER_VERSION=${FLASHINFER_VERSION} \ - SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=/opt/trtllm_gen_moe_cubin_pool + FLASHINFER_VERSION=${FLASHINFER_VERSION} # Add GKE default lib and bin locations ENV PATH="${PATH}:/usr/local/nvidia/bin" \ @@ -73,7 +69,6 @@ RUN --mount=type=cache,target=/var/cache/apt,id=base-apt \ build-essential \ cmake \ perl \ - patch \ patchelf \ ccache \ git-lfs \ @@ -164,7 +159,6 @@ ENV LANG=en_US.UTF-8 \ # | # +-- devtools_builder (independent) # +-- gateway_builder (independent, only needs gateway source) -# +-- trtllm_cubin_builder (independent) # | # v # framework (combines all artifacts) @@ -400,27 +394,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \ && cp target/release/sgl-model-gateway /build/sgl-model-gateway-bin \ && rm -rf /root/.cargo /root/.rustup /build/sgl-model-gateway/target /build/sgl-model-gateway/bindings/python/target -######################################################## -# PARALLEL STAGE 6: TRT-LLM Generated-MoE Cubin Pool -######################################################## -FROM base AS trtllm_cubin_builder - -RUN cubin_archive="/tmp/trtllm_gen_moe_cubin_pool.zip" && \ - cubin_extract_dir="/tmp/trtllm_gen_moe_cubin_extract" && \ - wget --no-verbose --output-document="${cubin_archive}" \ - "${TRTLLM_GEN_MOE_CUBIN_URL}" && \ - echo "${TRTLLM_GEN_MOE_CUBIN_SHA256} ${cubin_archive}" | \ - sha256sum --check --strict - && \ - mkdir -p "${cubin_extract_dir}" && \ - unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" && \ - test ! -e "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \ - mv "${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" \ - "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \ - test "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" \ - -type f -name '*.cubin' | wc -l)" -eq 1696 && \ - rm -f "${cubin_archive}" && \ - rm -rf "${cubin_extract_dir}" - ######################################################## ########## Final Framework Image ###################### ######################################################## @@ -455,21 +428,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \ # Copy flashinfer cubin (always) and jit-cache (if installed) packages COPY --from=flashinfer_cache /flashinfer_jit_output/ /usr/local/lib/python3.12/dist-packages/ -# Apply the FlashInfer CuTeDSL MLA decode-context-parallel runtime patch. -# Exclude tests because they are not included in the installed wheel. -COPY docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt /tmp/flashinfer-perkz-dcp-0.6.15.txt -RUN FLASHINFER_DCP_PATCH=/tmp/flashinfer-perkz-dcp-0.6.15.txt && \ - FLASHINFER_SITE_PACKAGES="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')" && \ - sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \ - patch --dry-run --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \ - sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \ - patch --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \ - rm -f "${FLASHINFER_DCP_PATCH}" && \ - rm -rf /root/.cache/flashinfer /root/.cache/pip - -# Copy the pinned FlashInfer MXFP4 MoE runner cubin pool -COPY --from=trtllm_cubin_builder /opt/trtllm_gen_moe_cubin_pool /opt/trtllm_gen_moe_cubin_pool - # Copy dev tools COPY --from=devtools_builder /tools/diff-so-fancy /usr/local/bin/ COPY --from=devtools_builder /tools/clang-format /usr/local/bin/ @@ -734,8 +692,7 @@ ARG GDRCOPY_VERSION=2.5.1 ENV DEBIAN_FRONTEND=noninteractive \ CUDA_HOME=/usr/local/cuda \ - GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ \ - SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=/opt/trtllm_gen_moe_cubin_pool + GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ # Add GKE default lib and bin locations + CUDA compiler paths for FlashInfer JIT ENV PATH="${PATH}:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/cuda/nvvm/bin" \ @@ -821,9 +778,6 @@ RUN --mount=type=cache,target=/var/cache/apt,id=runtime-apt \ # Copy Python site-packages from framework (already cleaned of __pycache__/tests/pyc files) COPY --from=framework_final /usr/local/lib/python3.12/dist-packages /usr/local/lib/python3.12/dist-packages -# Copy the pinned FlashInfer MXFP4 MoE runner cubin pool -COPY --from=framework_final /opt/trtllm_gen_moe_cubin_pool /opt/trtllm_gen_moe_cubin_pool - # Copy SGLang workspace COPY --from=framework_final /sgl-workspace /sgl-workspace diff --git a/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt b/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt deleted file mode 100644 index 945fbe395..000000000 --- a/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt +++ /dev/null @@ -1,5639 +0,0 @@ -diff --git a/flashinfer/cute_dsl/attention/mla_dispatch.py b/flashinfer/cute_dsl/attention/mla_dispatch.py -index f41976df..c56b21c5 100644 ---- a/flashinfer/cute_dsl/attention/mla_dispatch.py -+++ b/flashinfer/cute_dsl/attention/mla_dispatch.py -@@ -9,8 +9,10 @@ controlled by the ``cute_dsl_impl`` kwarg, with three valid values: - - * ``"auto"`` (default) — library picks the right implementation. - Monolithic by default, automatically promoted to modular when the call -- uses a feature monolithic doesn't support (currently: ``sinks``). -+ uses a feature monolithic doesn't support (currently: ``sinks``). DCP -+ requires monolithic and therefore keeps the monolithic implementation. - * ``"modular"`` — strict. Always run the modular implementation. -+ Raises :class:`ValueError` when DCP is enabled. - * ``"monolithic"`` — strict. Always run the monolithic implementation; - raise :class:`ValueError` if the call uses any modular-only feature. - No silent fallback — the contract is "you asked for monolithic, you -@@ -45,6 +47,12 @@ _logged_impls: set[str] = set() - # Add new entries here as more variants are exposed through the standalone - # signature. - MODULAR_ONLY_KWARGS = ("sinks",) -+DCP_KWARGS = ( -+ "enable_dcp", -+ "cp_world", -+ "cp_rank", -+ "causal_seqlens_kv_global", -+) - - - def _has_modular_only_feature(kwargs: dict) -> Optional[str]: -@@ -56,6 +64,45 @@ def _has_modular_only_feature(kwargs: dict) -> Optional[str]: - return None - - -+def _validate_dcp_kwargs(kwargs: dict) -> bool: -+ """Validate dispatcher-level DCP selection and return whether it is enabled. -+ -+ Tensor shape, dtype, and device validation belongs to the monolithic -+ wrapper, which has access to the normalized query batch. The dispatcher -+ only enforces the static feature contract and prevents non-default DCP -+ arguments from being silently ignored. -+ """ -+ enable_dcp = kwargs.get("enable_dcp", False) -+ if not isinstance(enable_dcp, bool): -+ raise TypeError(f"enable_dcp must be a bool, got {type(enable_dcp).__name__}") -+ -+ cp_world = kwargs.get("cp_world", 1) -+ cp_rank = kwargs.get("cp_rank", 0) -+ causal_seqlens_kv_global = kwargs.get("causal_seqlens_kv_global") -+ if not isinstance(cp_world, int) or isinstance(cp_world, bool) or cp_world <= 0: -+ raise ValueError(f"cp_world must be a positive integer, got {cp_world!r}") -+ if not isinstance(cp_rank, int) or isinstance(cp_rank, bool): -+ raise TypeError(f"cp_rank must be an integer, got {type(cp_rank).__name__}") -+ if enable_dcp and not 0 <= cp_rank < cp_world: -+ raise ValueError( -+ f"cp_rank must satisfy 0 <= cp_rank < cp_world, got " -+ f"cp_rank={cp_rank}, cp_world={cp_world}" -+ ) -+ if not enable_dcp: -+ nondefault = [] -+ if cp_world != 1: -+ nondefault.append(f"cp_world={cp_world!r}") -+ if cp_rank != 0: -+ nondefault.append(f"cp_rank={cp_rank!r}") -+ if causal_seqlens_kv_global is not None: -+ nondefault.append("causal_seqlens_kv_global") -+ if nondefault: -+ raise ValueError( -+ "DCP arguments require enable_dcp=True; got " + ", ".join(nondefault) -+ ) -+ return enable_dcp -+ -+ - def _resolve_impl(*, requested: str, kwargs: dict) -> str: - """Map a user request and call kwargs to a concrete impl name. - -@@ -67,9 +114,19 @@ def _resolve_impl(*, requested: str, kwargs: dict) -> str: - f"Invalid cute_dsl_impl={requested!r}; expected one of {VALID_IMPLS}" - ) - -+ enable_dcp = _validate_dcp_kwargs(kwargs) - needs_modular = _has_modular_only_feature(kwargs) - -+ if enable_dcp and needs_modular is not None: -+ raise ValueError( -+ f"DCP cannot be combined with {needs_modular!r}: DCP requires the " -+ "monolithic CuTeDSL MLA implementation, while the requested feature " -+ "requires the modular implementation." -+ ) -+ - if requested == "auto": -+ if enable_dcp: -+ return "monolithic" - return "modular" if needs_modular is not None else _DEFAULT_IMPL - - if requested == "monolithic" and needs_modular is not None: -@@ -80,6 +137,13 @@ def _resolve_impl(*, requested: str, kwargs: dict) -> str: - f"right impl based on the call) or cute_dsl_impl='modular'." - ) - -+ if requested == "modular" and enable_dcp: -+ raise ValueError( -+ "cute_dsl_impl='modular' was requested with enable_dcp=True, but " -+ "DCP is only supported by the monolithic CuTeDSL MLA implementation. " -+ "Use cute_dsl_impl='auto' or cute_dsl_impl='monolithic'." -+ ) -+ - return requested # "modular" or "monolithic" - - -@@ -91,18 +155,19 @@ def cute_dsl_mla_decode(*args, cute_dsl_impl: str = "auto", **kwargs): - :func:`flashinfer.cute_dsl.attention.wrappers.batch_mla.cute_dsl_mla_decode` - (modular, supports ``sinks=``) and - :func:`flashinfer.cute_dsl.attention.monolithic.mla_decode.cute_dsl_mla_decode` -- (monolithic, no variant support) — their signatures are otherwise -- identical. -+ (monolithic, supports static DCP). Implementation-specific keyword -+ arguments are validated before dispatch and stripped only when their -+ default values make them irrelevant to the selected implementation. - - Parameters - ---------- - cute_dsl_impl : str, default ``"auto"`` - ``"auto"`` (default) lets the dispatcher pick: monolithic by - default, modular when the call uses a modular-only feature -- (currently ``sinks``). ``"modular"`` and ``"monolithic"`` are -- strict — the dispatcher will not silently switch implementations, -- and ``"monolithic"`` raises :class:`ValueError` if the call uses -- a modular-only feature. -+ (currently ``sinks``); DCP requires monolithic. ``"modular"`` and -+ ``"monolithic"`` are strict — the dispatcher will not silently switch -+ implementations, and raises :class:`ValueError` for an incompatible -+ feature request. - """ - impl = _resolve_impl(requested=cute_dsl_impl, kwargs=kwargs) - -@@ -122,5 +187,8 @@ def cute_dsl_mla_decode(*args, cute_dsl_impl: str = "auto", **kwargs): - kwargs = {k: v for k, v in kwargs.items() if k not in MODULAR_ONLY_KWARGS} - from .monolithic.mla_decode import cute_dsl_mla_decode as _impl - else: -+ # Default-valued DCP kwargs are accepted by the common dispatcher but -+ # are not part of the modular implementation's ABI. -+ kwargs = {k: v for k, v in kwargs.items() if k not in DCP_KWARGS} - from .wrappers.batch_mla import cute_dsl_mla_decode as _impl - return _impl(*args, **kwargs) -diff --git a/flashinfer/cute_dsl/attention/monolithic/mla_decode.py b/flashinfer/cute_dsl/attention/monolithic/mla_decode.py -index cb491ae1..eb9be701 100644 ---- a/flashinfer/cute_dsl/attention/monolithic/mla_decode.py -+++ b/flashinfer/cute_dsl/attention/monolithic/mla_decode.py -@@ -32,6 +32,7 @@ from flashinfer.utils import device_support_pdl - - from .mla_decode_fp16 import BlackwellMultiHeadLatentAttentionForwardFP16 - from .mla_decode_fp8 import BlackwellMultiHeadLatentAttentionForwardFP8 -+from .mla_helpers import MAX_SPLITS, ceil_div, compute_q_tile_layout - from flashinfer.cute_dsl.utils import ( - _as_cute_dsl_workspace_i8, - get_max_active_clusters, -@@ -40,6 +41,60 @@ from flashinfer.cute_dsl.utils import ( - ) - - -+_CUDA_GRID_Y_MAX = 65_535 -+_REDUCER_D_TILE_CANDIDATES = (1, 2, 4) -+_STATIC_REDUCER_MAX_SPLITS = 32 -+ -+ -+def _validate_nonpersistent_grid_y( -+ batch_size: int, num_q_tiles: int, is_persistent: bool -+) -> None: -+ """Reject nonpersistent launch grids whose Y dimension CUDA cannot encode.""" -+ grid_y = batch_size * num_q_tiles -+ if not is_persistent and grid_y > _CUDA_GRID_Y_MAX: -+ raise ValueError( -+ "nonpersistent CuTeDSL MLA grid.y would be " -+ f"{grid_y}, exceeding the CUDA limit {_CUDA_GRID_Y_MAX} " -+ f"(batch_size={batch_size}, num_q_tiles={num_q_tiles})" -+ ) -+ -+ -+def _get_reducer_d_tiles( -+ batch_size: int, -+ seq_len_q: int, -+ num_heads: int, -+ num_sms: int, -+ effective_split_kv: int, -+) -> int: -+ """Choose output-side reducer parallelism for an underfilled row grid. -+ -+ A reducer CTA normally owns one D512 row. When the real row grid does not -+ cover one resident wave, compare the conservative wave count for one, two, -+ or four equal D bands and keep the smallest topology with the shortest -+ per-band critical path. Once rows already cover every SM, retain one CTA -+ per row and avoid duplicated LSE work. -+ """ -+ reducer_rows = batch_size * seq_len_q * num_heads -+ if ( -+ reducer_rows <= 0 -+ or num_sms <= 0 -+ or reducer_rows >= num_sms -+ or effective_split_kv <= 1 -+ ): -+ return 1 -+ -+ best_tiles = 1 -+ best_waves = ceil_div(reducer_rows, num_sms) -+ for d_tiles in _REDUCER_D_TILE_CANDIDATES[1:]: -+ if d_tiles > effective_split_kv: -+ continue -+ waves = ceil_div(reducer_rows * d_tiles, num_sms) -+ if waves * best_tiles < best_waves * d_tiles: -+ best_tiles = d_tiles -+ best_waves = waves -+ return best_tiles -+ -+ - @functools.cache - def _get_split_kv_and_workspace_size( - B: int, -@@ -47,22 +102,32 @@ def _get_split_kv_and_workspace_size( - H: int, - kv_lora_rank: int, - max_active_blocks: int, -+ max_seq_len: Optional[int] = None, - ) -> Tuple[int, int]: -- """Cache split_kv and workspace_size since they are deterministic for the same params.""" -- # When folding S_q into heads, the workspace dims are the effective dims -- # (num_heads * F, q_len // F). get_workspace_size already pads H<128 to -- # 128, so passing num_heads_eff and seq_len_q_eff yields the right size. -+ """Return the nonempty split count and its workspace requirement. -+ -+ The occupancy heuristic can request more splits than the kernel's uniform -+ contiguous partitioning actually uses. When ``max_seq_len`` is known, -+ normalize the candidate to the number of nonempty fixed-size K chunks so -+ the grid and reducer do not carry an empty final split. -+ """ -+ # Flatten (query_token, head) into one contiguous row space. A cooperative -+ # 2-CTA MMA tile owns 128 consecutive rows and may cross token boundaries; -+ # only the final query tile can be padded. - mma_qk_tile_m = 128 -- fold_sq_ratio = BlackwellMultiHeadLatentAttentionForwardFP16.compute_fold_sq_ratio( -- H, q_len, mma_qk_tile_m -- ) -- num_heads_eff = H * fold_sq_ratio -- seq_len_q_eff = q_len // fold_sq_ratio -+ mma_qk_tile_n = 128 -+ _, num_q_tiles, _ = compute_q_tile_layout(H, q_len, mma_qk_tile_m) - split_kv = BlackwellMultiHeadLatentAttentionForwardFP16.get_split_kv_simplified( -- B, seq_len_q_eff, max_active_blocks -+ B, num_q_tiles, max_active_blocks - ) -+ if max_seq_len is not None: -+ if max_seq_len <= 0: -+ raise ValueError(f"max_seq_len must be > 0, got {max_seq_len}") -+ k_tile_total = ceil_div(max_seq_len, mma_qk_tile_n) -+ k_tiles_per_split = ceil_div(k_tile_total, split_kv) -+ split_kv = ceil_div(k_tile_total, k_tiles_per_split) - workspace_size = BlackwellMultiHeadLatentAttentionForwardFP16.get_workspace_size( -- num_heads_eff, seq_len_q_eff, kv_lora_rank, B, split_kv, cutlass.Float32 -+ mma_qk_tile_m, num_q_tiles, kv_lora_rank, B, split_kv, cutlass.Float32 - ) - return split_kv, workspace_size - -@@ -79,8 +144,19 @@ def _check_can_implement( - is_persistent: bool, - is_var_seq: bool, - is_var_split_kv: bool, -+ enable_dcp: bool = False, -+ cp_world: int = 1, - ) -> None: - """Check if the kernel supports the given configuration (cached).""" -+ if not isinstance(enable_dcp, bool): -+ raise TypeError(f"enable_dcp must be a bool, got {type(enable_dcp).__name__}") -+ if not isinstance(cp_world, int) or isinstance(cp_world, bool) or cp_world <= 0: -+ raise ValueError(f"cp_world must be a positive integer, got {cp_world!r}") -+ if not enable_dcp and cp_world != 1: -+ raise ValueError( -+ f"cp_world={cp_world} requires enable_dcp=True; disabled DCP uses cp_world=1" -+ ) -+ - mma_qk_tiler_mn = (128, 128) - mma_pv_tiler_mn = (128, 256) - -@@ -129,15 +205,20 @@ def _get_compiled_mla_kernel( - is_persistent: bool, - is_var_seq: bool, - is_var_split_kv: bool, -+ reducer_d_tiles: int = 1, -+ reducer_max_splits: int = MAX_SPLITS, - skip_correction_threshold: float = 0.0, - is_workspace_size_zero: bool = False, - enable_pdl: bool = False, -+ enable_dcp: bool = False, -+ cp_world: int = 1, - ) -> Callable: - """Compile and cache an MLA decode kernel. - - Returns a callable that accepts (q_latent, q_rope, c_latent, c_rope, - page_table, o, lse, workspace, split_kv_scalar, cache_seqs, -- block_split_kvs, softmax_scale_scalar, output_scale_scalar). -+ causal_seqlens_kv_global, cp_rank_scalar, block_split_kvs, -+ softmax_scale_scalar, output_scale_scalar). - - All scalar arguments must be pre-wrapped as Int32/Float32. - """ -@@ -157,14 +238,6 @@ def _get_compiled_mla_kernel( - cutlass_dtype = torch_to_cutlass_dtype(torch_dtype) - cutlass_out_dtype = torch_to_cutlass_dtype(torch_out_dtype) - -- # Derive the seq_len_q-into-heads fold factor. F > 1 means the kernel -- # repacks the [H, S_q] tile to [H*F, S_q/F] internally so MTP / spec-decoding -- # with H < 128 fully populates the 128-wide MMA-M tile. -- fold_sq_ratio = KernelClass.compute_fold_sq_ratio( -- num_heads, seq_len_q, mma_qk_tiler_mn[0] -- ) -- fold_sq = fold_sq_ratio > 1 -- - kernel_obj = KernelClass( - acc_dtype=cutlass.Float32, - lse_dtype=cutlass.Float32, -@@ -181,7 +254,10 @@ def _get_compiled_mla_kernel( - enable_pdl=enable_pdl, - num_heads=num_heads, - seq_len_q=seq_len_q, -- fold_sq=fold_sq, -+ reducer_d_tiles=reducer_d_tiles, -+ reducer_max_splits=reducer_max_splits, -+ enable_dcp=enable_dcp, -+ cp_world=cp_world, - ) - - # All dimensions as sym_int — this matches the original kernel's use of -@@ -271,6 +347,18 @@ def _get_compiled_mla_kernel( - (sym_batch,), - assumed_align=16, - ) -+ # DCP's causal boundary is global while cache_seqs remains the physical -+ # rank-local length used for paging and split-K traversal. Keep the tensor -+ # absent from the disabled specialization so no load or pointer argument -+ # reaches its generated device code. -+ if enable_dcp: -+ causal_seqlens_kv_global_fake = cute.runtime.make_fake_compact_tensor( -+ cutlass.Int32, -+ (sym_batch,), -+ assumed_align=4, -+ ) -+ else: -+ causal_seqlens_kv_global_fake = None - # block_split_kvs: [batch_size] — int32 (only needed for is_var_split_kv=True) - if is_var_split_kv: - block_split_kvs_fake = cute.runtime.make_fake_compact_tensor( -@@ -295,6 +383,8 @@ def _get_compiled_mla_kernel( - workspace_fake, - Int32(1), # split_kv placeholder - cache_seqs_fake, -+ causal_seqlens_kv_global_fake, -+ Int32(0), # cp_rank placeholder (runtime-uniform, not a JIT key) - block_split_kvs_fake, - Float32(1.0), # softmax_scale placeholder - Float32(1.0), # output_scale placeholder -@@ -323,6 +413,10 @@ def cute_dsl_mla_decode( - enable_pdl: Optional[bool] = None, - lse: Optional[torch.Tensor] = None, - return_lse: bool = False, -+ enable_dcp: bool = False, -+ cp_world: int = 1, -+ cp_rank: int = 0, -+ causal_seqlens_kv_global: Optional[torch.Tensor] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - """CuTe DSL MLA decode kernel for Blackwell SM100. - -@@ -336,8 +430,9 @@ def cute_dsl_mla_decode( - Pre-allocated workspace buffer (int8 or uint8). Required size depends on batch size - and split_kv (auto-computed from B, q_len, and number of SMs): - -- - Formula: ``B * H * q_len * split_kv * (kv_lora_rank + 1) * 4`` bytes -- (0 when split_kv == 1, which happens when B >= num_SMs / 2) -+ - Formula: ``B * 128 * q_tiles * split_kv * (kv_lora_rank + 1) * 4`` -+ bytes, where ``q_tiles = ceil(q_len * H / 128)`` -+ (0 when ``split_kv == 1``, typically once ``B * q_tiles >= num_SMs / 2``) - - Typical max: ~18 MB on a 148-SM GPU (e.g. B=4..8, H=128, D=512) - - Safe default: 128 MB covers all realistic configurations - kv_lora_rank : int -@@ -355,7 +450,7 @@ def cute_dsl_mla_decode( - output_scale : float - Scale factor applied to the output. - out : Optional[torch.Tensor] -- Pre-allocated output tensor [B, q_len, H, kv_lora_rank]. -+ Contiguous pre-allocated output tensor [B, q_len, H, kv_lora_rank]. - out_dtype : Optional[torch.dtype] - Output data type. If None, defaults to torch.bfloat16 (matching trtllm-gen backend). - Supported values: torch.bfloat16, torch.float8_e4m3fn (FP8 input only), -@@ -375,13 +470,25 @@ def cute_dsl_mla_decode( - * ``[B * q_len, H]`` (matches ``trtllm-gen`` shape; the wrapper - reshapes via ``.view`` to the native layout). - -+ Caller-provided buffers must be contiguous. -+ - If ``return_lse`` is True and this is None, a buffer of the native - ``[B, q_len, H]`` shape is allocated internally. - return_lse : bool - Whether to return LSE values. When True, the function returns - ``(out, lse)`` (the ``lse`` tensor returned is in whatever shape - the caller supplied; if no ``lse`` was supplied, ``[B, q_len, H]``). -- -+ enable_dcp : bool -+ Enable static cyclic decode context-parallel masking. DCP returns a -+ rank-local attention state and therefore requires ``return_lse=True``. -+ cp_world : int -+ Compile-time context-parallel world size. The local cache owns global -+ token positions ``cp_world * local_k + cp_rank``. -+ cp_rank : int -+ Runtime-uniform context-parallel rank. -+ causal_seqlens_kv_global : Optional[torch.Tensor] -+ Contiguous CUDA int32 tensor ``[B]`` containing the global exclusive -+ causal bound for the newest query token. Required when DCP is enabled. - Returns - ------- - torch.Tensor or Tuple[torch.Tensor, torch.Tensor] -@@ -395,9 +502,88 @@ def cute_dsl_mla_decode( - assert kv_cache.dtype == query.dtype, ( - f"kv_cache dtype {kv_cache.dtype} must match query dtype {query.dtype}" - ) -+ if query.ndim != 4: -+ raise ValueError( -+ "query must have shape " -+ "[batch_size, q_len_per_request, num_heads, head_dim_qk]" -+ ) - B, q_len, H, D_qk = query.shape - assert D_qk == kv_lora_rank + qk_rope_head_dim - -+ if not isinstance(enable_dcp, bool): -+ raise TypeError(f"enable_dcp must be a bool, got {type(enable_dcp).__name__}") -+ if not isinstance(cp_world, int) or isinstance(cp_world, bool) or cp_world <= 0: -+ raise ValueError(f"cp_world must be a positive integer, got {cp_world!r}") -+ if not isinstance(cp_rank, int) or isinstance(cp_rank, bool): -+ raise TypeError(f"cp_rank must be an integer, got {type(cp_rank).__name__}") -+ -+ if enable_dcp: -+ if not 0 <= cp_rank < cp_world: -+ raise ValueError( -+ f"cp_rank must satisfy 0 <= cp_rank < cp_world, got " -+ f"cp_rank={cp_rank}, cp_world={cp_world}" -+ ) -+ if not return_lse: -+ raise ValueError( -+ "enable_dcp=True requires return_lse=True so rank-local states " -+ "can be merged with their LSE values" -+ ) -+ if causal_seqlens_kv_global is None: -+ raise ValueError( -+ "causal_seqlens_kv_global is required when enable_dcp=True" -+ ) -+ if not isinstance(causal_seqlens_kv_global, torch.Tensor): -+ raise TypeError( -+ "causal_seqlens_kv_global must be a torch.Tensor, got " -+ f"{type(causal_seqlens_kv_global).__name__}" -+ ) -+ if causal_seqlens_kv_global.dtype != torch.int32: -+ raise ValueError( -+ "causal_seqlens_kv_global must have dtype torch.int32, got " -+ f"{causal_seqlens_kv_global.dtype}" -+ ) -+ if not causal_seqlens_kv_global.is_cuda: -+ raise ValueError("causal_seqlens_kv_global must be a CUDA tensor") -+ if causal_seqlens_kv_global.device != query.device: -+ raise ValueError( -+ "causal_seqlens_kv_global must be on the query device " -+ f"{query.device}, got {causal_seqlens_kv_global.device}" -+ ) -+ if tuple(causal_seqlens_kv_global.shape) != (B,): -+ raise ValueError( -+ f"causal_seqlens_kv_global must have shape ({B},), got " -+ f"{tuple(causal_seqlens_kv_global.shape)}" -+ ) -+ if not causal_seqlens_kv_global.is_contiguous(): -+ raise ValueError("causal_seqlens_kv_global must be contiguous") -+ else: -+ nondefault = [] -+ if cp_world != 1: -+ nondefault.append(f"cp_world={cp_world}") -+ if cp_rank != 0: -+ nondefault.append(f"cp_rank={cp_rank}") -+ if causal_seqlens_kv_global is not None: -+ nondefault.append("causal_seqlens_kv_global") -+ if nondefault: -+ raise ValueError( -+ "DCP arguments require enable_dcp=True; got " + ", ".join(nondefault) -+ ) -+ -+ # The single Q TMA descriptor flattens (query_token, head) into one -+ # globally bounded affine row mode. Adjacent tokens must therefore follow -+ # all H head rows. Normal contiguous queries and last-dimension slices -+ # satisfy this; materialize uncommon token-strided/head-subset views. -+ if query.stride(-1) != 1 or query.stride(1) != H * query.stride(2): -+ query = query.contiguous() -+ -+ # O/LSE are compiled with compact layouts. In particular, a token-gapped -+ # view cannot be flattened across an M128 tile boundary and would be -+ # silently written through the gap, so reject noncompact buffers early. -+ if out is not None and not out.is_contiguous(): -+ raise ValueError(f"out must be contiguous, got strides {out.stride()}") -+ if lse is not None and not lse.is_contiguous(): -+ raise ValueError(f"lse must be contiguous, got strides {lse.stride()}") -+ - q_dtype = query.dtype - # Resolve output dtype: for FP8 input, default to bfloat16 (matching trtllm-gen backend); - # for FP16/BF16 input, default to same as input. Allow override via out_dtype or out tensor. -@@ -415,7 +601,8 @@ def cute_dsl_mla_decode( - kv_cache = kv_cache.squeeze(1) - page_size = kv_cache.shape[1] - -- # Split query into latent and rope components — keep contiguous [B, q_len, H, D]. -+ # Split query into latent and rope components. The slices share the flat -+ # affine token/head row stride established above. - # The kernel's __call__ reinterprets to [H, D, q_len, B] via zero-cost make_tensor. - q_latent_k = query[..., :kv_lora_rank] - q_rope_k = query[..., kv_lora_rank:] -@@ -434,9 +621,13 @@ def cute_dsl_mla_decode( - # Cached split_kv and workspace_size computation - max_active_blocks = get_num_sm(query.device) - split_kv, workspace_size = _get_split_kv_and_workspace_size( -- B, q_len, H, kv_lora_rank, max_active_blocks -+ B, q_len, H, kv_lora_rank, max_active_blocks, max_seq_len - ) - -+ _, num_q_tiles, _ = compute_q_tile_layout(H, q_len) -+ is_persistent = not is_var_seq -+ _validate_nonpersistent_grid_y(B, num_q_tiles, is_persistent) -+ - # Prepare workspace: the CUTE signature uses int8, while public FlashInfer - # workspace buffers are byte storage and may be uint8. - workspace_buffer = _as_cute_dsl_workspace_i8(workspace_buffer) -@@ -446,6 +637,14 @@ def cute_dsl_mla_decode( - f"need {workspace_size} bytes" - ) - is_workspace_size_zero = workspace_size == 0 -+ k_tile_total = ceil_div(max_seq_len, 128) -+ k_tiles_per_split = ceil_div(k_tile_total, split_kv) -+ effective_split_kv = ceil_div(k_tile_total, k_tiles_per_split) -+ reducer_d_tiles = ( -+ 1 -+ if is_workspace_size_zero -+ else _get_reducer_d_tiles(B, q_len, H, max_active_blocks, effective_split_kv) -+ ) - if is_workspace_size_zero: - workspace_bytes = None - else: -@@ -486,9 +685,6 @@ def cute_dsl_mla_decode( - block_split_kvs = None - skip_correction_threshold = 0.0 - -- # for fix-length, set is_persistent to True; otherwise, set to False. -- is_persistent = not is_var_seq -- - # Validate configuration (cached, negligible overhead after first call) - _check_can_implement( - torch_dtype=q_dtype, -@@ -501,6 +697,8 @@ def cute_dsl_mla_decode( - is_persistent=is_persistent, - is_var_seq=is_var_seq, - is_var_split_kv=is_var_split_kv, -+ enable_dcp=enable_dcp, -+ cp_world=cp_world, - ) - - enable_pdl = device_support_pdl(query.device) if enable_pdl is None else enable_pdl -@@ -519,9 +717,13 @@ def cute_dsl_mla_decode( - is_persistent=is_persistent, - is_var_seq=is_var_seq, - is_var_split_kv=is_var_split_kv, -+ reducer_d_tiles=reducer_d_tiles, -+ reducer_max_splits=_STATIC_REDUCER_MAX_SPLITS, - skip_correction_threshold=skip_correction_threshold, - is_workspace_size_zero=is_workspace_size_zero, - enable_pdl=enable_pdl, -+ enable_dcp=enable_dcp, -+ cp_world=cp_world, - ) - - # Call the kernel -@@ -536,6 +738,8 @@ def cute_dsl_mla_decode( - workspace_bytes, - Int32(split_kv), - cache_seqs, -+ causal_seqlens_kv_global, -+ Int32(cp_rank), - block_split_kvs, - Float32(softmax_scale), - Float32(output_scale), -diff --git a/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp16.py b/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp16.py -index 8e33acee..8202fb3b 100644 ---- a/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp16.py -+++ b/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp16.py -@@ -77,6 +77,7 @@ from cutlass.cutlass_dsl import BaseDSL - - from .mla_helpers import ( - ceil_div, -+ compute_q_tile_layout, - MAX_SPLITS, - LOG2_E, - MLAStaticTileScheduler, -@@ -108,15 +109,17 @@ launcher and ``flashinfer/cute_dsl/attention/mla_dispatch.py`` for impl selectio - - Constraints: - * Data type requirements: -- - Input/output: Float16 -+ - Input: Float16 or BFloat16 -+ - Output: Float16 or BFloat16 - - Accumulation and LSE: Float32 - * Fixed architecture parameters: -- - Number of attention heads: 128 -+ - Number of attention heads: 1-128 - - Latent dimension: 512 - - RoPE dimension: 64 - * Input query modes should be (NumHeads, LatentDim/RopeDim, SeqLenQ, BatchSize) - * Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize) --* Query sequence length must be 1-4 -+* Query sequence length must be positive; query (token, head) rows are packed -+ continuously into 128-row M tiles with a safely padded final tile - * Only supports 2-CTA instructions - * Variable sequence length requires page table storage enabled - """ -@@ -138,7 +141,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - enable_pdl: bool, - num_heads: int = 128, - seq_len_q: int = 1, -- fold_sq: bool = False, -+ reducer_d_tiles: int = 1, -+ reducer_max_splits: int = MAX_SPLITS, -+ enable_dcp: bool = False, -+ cp_world: int = 1, - ): - """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel. - -@@ -164,17 +170,24 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - :type is_var_split_kv: bool - :param enable_pdl: Whether to use PDL - :type enable_pdl: bool -- :param num_heads: Number of attention heads (pre-fold). Used for the -- per-row spec-decoding (MTP) causal mask q_token_index computation. -+ :param num_heads: Number of attention heads. Defines the flattened -+ ``(query_token, head)`` row geometry and division-free causal mask. - :type num_heads: int -- :param seq_len_q: Query sequence length (pre-fold). Combined with -- ``num_heads`` to derive the per-row q_token used by the causal mask. -+ :param seq_len_q: Query sequence length. Combined with ``num_heads`` -+ to define the flattened query-row extent and causal boundary. - :type seq_len_q: int -- :param fold_sq: Whether to fold tokens of ``seq_len_q`` into the head -- dimension so the M tile becomes [F sub_q_tok][num_heads heads]. -- Required when ``num_heads < mma_qk_tiler_mn[0]`` and ``seq_len_q > 1`` -- so the M tile is fully populated. -- :type fold_sq: bool -+ :param reducer_d_tiles: Number of independent D bands reduced per row. -+ :type reducer_d_tiles: int -+ :param reducer_max_splits: Compile-time reducer capacity. Direct users -+ retain the generic 256-split capacity by default; callers choosing -+ a smaller specialization must cap runtime split-KV accordingly. -+ :type reducer_max_splits: int -+ :param enable_dcp: Whether to compile the decode-context-parallel -+ global-coordinate causal mask. -+ :type enable_dcp: bool -+ :param cp_world: Number of cyclic context-parallel shards. This is a -+ compile-time parameter when DCP is enabled. -+ :type cp_world: int - """ - - self.latent_dim = 512 -@@ -189,25 +202,37 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - self.page_size = page_size - self.is_var_seq = is_var_seq - self.is_var_split_kv = is_var_split_kv -+ if not 1 <= reducer_max_splits <= MAX_SPLITS: -+ raise ValueError( -+ f"reducer_max_splits must be in [1, {MAX_SPLITS}], " -+ f"got {reducer_max_splits}" -+ ) -+ if is_var_split_kv and reducer_max_splits != MAX_SPLITS: -+ raise ValueError( -+ "variable split-KV requires the generic reducer capacity " -+ f"of {MAX_SPLITS}, got {reducer_max_splits}" -+ ) -+ self.reducer_max_splits = reducer_max_splits - self.enable_pdl = enable_pdl -- # Original (pre-fold) num_heads and seq_len_q used for per-row -- # spec-decoding (MTP) causal q_token_index computation. When fold_sq is -- # True the M tile is laid out as [F sub_q_tok][num_heads heads]; the -- # full q_tok for row r is blk_coord[1] * F + (r // num_heads). -+ if cp_world < 1: -+ raise ValueError(f"cp_world must be positive, got {cp_world}") -+ if not enable_dcp and cp_world != 1: -+ raise ValueError( -+ "cp_world must be 1 when DCP is disabled, " -+ f"got cp_world={cp_world}" -+ ) -+ self.enable_dcp = enable_dcp -+ self.cp_world = cp_world -+ # Flatten (query token, head) into the physical 128-wide MMA M mode. -+ # Token boundaries may cross a tile; only the final global tile can be -+ # partially populated and is padded by the Q TMA OOB fill. - self.num_heads = num_heads - self.seq_len_q = seq_len_q -- # fold_sq (caller-controlled): whether the folding code path is enabled. -- # fold_sq_ratio (derived): fold factor F ≥ 1; the largest divisor of -- # seq_len_q with num_heads * F ≤ M_tile and F ≤ seq_len_q. When the -- # caller passes fold_sq=False, the kernel ignores the ratio. -- # When fold_sq=True but the derived ratio is 1, the folding branch -- # is taken with F=1 (a no-op transform). -- self.fold_sq = fold_sq -- self.fold_sq_ratio = ( -- BlackwellMultiHeadLatentAttentionForwardFP16.compute_fold_sq_ratio( -- num_heads, seq_len_q, mma_qk_tiler_mn[0] -- ) -- ) -+ ( -+ self.total_q_rows, -+ self.num_q_tiles, -+ self.tail_q_rows, -+ ) = compute_q_tile_layout(num_heads, seq_len_q, mma_qk_tiler_mn[0]) - self.cluster_shape_mnk = (2, 1, 1) - self.use_2cta_instrs = True - # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2), -@@ -215,6 +240,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - self.warps_in_n = 2 - self.num_compute_warps = 4 - self.threads_per_warp = 32 -+ if reducer_d_tiles not in (1, 2, 4): -+ raise ValueError(f"unsupported reducer_d_tiles={reducer_d_tiles}") -+ self.reducer_d_tiles = reducer_d_tiles -+ self.reducer_d_tile = self.latent_dim // reducer_d_tiles - mma_qk_tiler_k = self.rope_dim if self.seq_len_q == 1 else self.rope_dim * 2 - self.mma_qk_tiler = ( - self.mma_qk_tiler_mn[0], -@@ -311,6 +340,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - workspace: cute.Tensor, - split_kv: cutlass.Int32, - cache_seqs: Optional[cute.Tensor], -+ causal_seqlens_kv_global: Optional[cute.Tensor], -+ cp_rank: cutlass.Int32, - block_split_kvs: Optional[cute.Tensor], - softmax_scale: cutlass.Float32, - output_scale: cutlass.Float32, -@@ -346,6 +377,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - :type split_kv: cutlass.Int32 - :param cache_seqs: The cache sequences tensor with shape [batch_size] - :type cache_seqs: cute.Tensor -+ :param causal_seqlens_kv_global: Global exclusive causal KV bound for -+ the newest query on each batch item. Present only for DCP kernels. -+ :type causal_seqlens_kv_global: cute.Tensor -+ :param cp_rank: Runtime context-parallel rank. -+ :type cp_rank: cutlass.Int32 - :param block_split_kvs: The block split KV tensor with shape [batch_size] - :type block_split_kvs: cute.Tensor - :param softmax_scale: The scale factor for softmax -@@ -386,6 +422,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - q_latent = _reinterpret_4d(q_latent) - q_rope = _reinterpret_4d(q_rope) - o = _reinterpret_4d(o) -+ # Keep unpacked output aliases for the split-KV reducer, which -+ # launches only real (head, q_token) rows. -+ o_unpacked = o - - # Reinterpret contiguous [num_pages, page_size, D] as [page_size, D, num_pages] - # Input stride: (PS*D, D, 1) → Target: (D, 1, PS*D) -@@ -419,54 +458,62 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - stride=(lse.stride[2], lse.stride[1], lse.stride[0]), - ), - ) -+ lse_unpacked = lse - -- # When num_heads < M tile, fold up to F = fold_sq_ratio tokens of -- # seq_len_q into the head dimension so M_eff = num_heads * F (≤ M_tile). -- # E.g., H=32, S_q=4 → F=4, M_eff=128, S_q_eff=1 -- # E.g., H=32, S_q=8 → F=4, M_eff=128, S_q_eff=2 -- # This works because MLA shares KV across all heads/queries independently. -- # Tensor layout: [H, D, S_q, B] → [H*F, D, S_q/F, B]; relies on -- # stride_S == stride_H * H (always true for contiguous [B, S_q, H, D] -- # tensors after _reinterpret_4d). -- if cutlass.const_expr(self.fold_sq): -- F = self.fold_sq_ratio -- -- def _fold_sq_4d(t): -- return cute.make_tensor( -- t.iterator, -- cute.make_layout( -- ( -- t.shape[0] * F, -- t.shape[1], -- t.shape[2] // F, -- t.shape[3], -- ), -- stride=( -- t.stride[0], -- t.stride[1], -- t.stride[2] * F, -- t.stride[3], -- ), -- ), -- ) -- -- q_latent = _fold_sq_4d(q_latent) -- q_rope = _fold_sq_4d(q_rope) -- o = _fold_sq_4d(o) -- # LSE: [H, S_q, B] → [H*F, S_q/F, B] -- lse = cute.make_tensor( -- lse.iterator, -+ # Flatten token-major [H, D, S_q, B] query storage into one globally -+ # bounded [S_q * H, D, B] M mode. Consecutive M128 tiles can cross -+ # token boundaries naturally, while TMA zero-fills only the final tail. -+ def _flatten_q_rows(t): -+ return cute.make_tensor( -+ t.iterator, - cute.make_layout( -- (lse.shape[0] * F, lse.shape[1] // F, lse.shape[2]), -- stride=(lse.stride[0], lse.stride[1] * F, lse.stride[2]), -+ (self.total_q_rows, t.shape[1], t.shape[3]), -+ stride=(t.stride[0], t.stride[1], t.stride[3]), - ), - ) - -+ q_latent = _flatten_q_rows(q_latent) -+ q_rope = _flatten_q_rows(q_rope) -+ -+ # O/LSE expose physical M128 query tiles to the main epilogue. Their -+ # final virtual padding is never dereferenced because stores are -+ # predicated by get_valid_q_rows(). -+ m_tile = self.mma_qk_tiler_mn[0] -+ runtime_num_q_tiles = cute.ceil_div(o.shape[0] * o.shape[2], m_tile) -+ o = cute.make_tensor( -+ o.iterator, -+ cute.make_layout( -+ ( -+ m_tile, -+ o.shape[1], -+ runtime_num_q_tiles, -+ o.shape[3], -+ ), -+ stride=( -+ o.stride[0], -+ o.stride[1], -+ o.stride[0] * m_tile, -+ o.stride[3], -+ ), -+ ), -+ ) -+ lse = cute.make_tensor( -+ lse.iterator, -+ cute.make_layout( -+ (m_tile, runtime_num_q_tiles, lse.shape[2]), -+ stride=( -+ lse.stride[0], -+ lse.stride[0] * m_tile, -+ lse.stride[2], -+ ), -+ ), -+ ) -+ - acc_o, acc_lse = self.initialize_workspace( -- q_latent.shape[0], -- q_latent.shape[1], -- q_latent.shape[2], -- q_latent.shape[3], -+ o.shape[0], -+ o.shape[1], -+ o.shape[2], -+ o.shape[3], - split_kv, - self.acc_dtype, - workspace, -@@ -732,6 +779,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - acc_lse, - split_kv, - cache_seqs, -+ causal_seqlens_kv_global, -+ cp_rank, - block_split_kvs, - softmax_scale_log2, - output_scale, -@@ -756,17 +805,21 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - ) - if cutlass.const_expr(acc_o is not None): - self.reduction_kernel( -- o, -- lse, -+ o_unpacked, -+ lse_unpacked, - acc_o, - acc_lse, - split_kv, - cache_seqs, - block_split_kvs, - ).launch( -- grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]), -+ grid=( -+ o_unpacked.shape[0] * self.reducer_d_tiles, -+ o_unpacked.shape[2], -+ o_unpacked.shape[3], -+ ), - block=[self.threads_per_warp * self.num_compute_warps, 1, 1], -- smem=MAX_SPLITS * self.acc_dtype.width // 8, -+ smem=self.reducer_max_splits * self.acc_dtype.width // 8, - stream=stream, - min_blocks_per_mp=1, - use_pdl=self.enable_pdl, -@@ -832,6 +885,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - mAccLSE: Optional[cute.Tensor], - split_kv: cutlass.Int32, - cache_seqs: cute.Tensor, -+ causal_seqlens_kv_global: Optional[cute.Tensor], -+ cp_rank: cutlass.Int32, - block_split_kvs: cute.Tensor, - softmax_scale_log2: cutlass.Float32, - output_scale: cutlass.Float32, -@@ -872,8 +927,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - :type mQL: cute.Tensor - :param tma_atom_q_rope: TMA copy atom for query rope tensor - :type tma_atom_q_rope: cute.CopyAtom -- :param mKR: Compressed rope tensor -- :type mKR: cute.Tensor -+ :param mQR: query rope tensor -+ :type mQR: cute.Tensor - :param tma_atom_c_latent: TMA copy atom for c latent tensor - :type tma_atom_c_latent: cute.CopyAtom - :param mCL: Compressed latent tensor -@@ -896,6 +951,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - :type split_kv: cutlass.Int32 - :param cache_seqs: The variable sequence length tensor - :type cache_seqs: cute.Tensor -+ :param causal_seqlens_kv_global: Per-batch global exclusive causal KV -+ bound for the newest query, present only when DCP is enabled. -+ :type causal_seqlens_kv_global: cute.Tensor -+ :param cp_rank: Runtime context-parallel rank. -+ :type cp_rank: cutlass.Int32 - :param block_split_kvs: The per-block split_kv values tensor - :type block_split_kvs: cute.Tensor - :param softmax_scale_log2: The log2 scale factor for softmax -@@ -1273,6 +1333,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - split_kv, cache_seqs, block_split_kvs, blk_coord - ) - if k_tile_count > 0: -+ causal_global = cutlass.Int32(0) -+ if cutlass.const_expr(self.enable_dcp): -+ causal_global = causal_seqlens_kv_global[blk_coord[2]] - compute_common_params = SimpleNamespace( - blk_coord=blk_coord, - split_kv=split_kv, -@@ -1281,6 +1344,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - mAccO=mAccO, - mO=mO, - K=cache_seqs[blk_coord[2]], -+ causal_global=causal_global, -+ cp_rank=cp_rank, - L=mCL.shape[1], - tmem_ptr=tmem_ptr, - tidx=tidx, -@@ -1337,6 +1402,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - split_kv, cache_seqs, block_split_kvs, blk_coord - ) - if k_tile_count > 0: -+ causal_global = cutlass.Int32(0) -+ if cutlass.const_expr(self.enable_dcp): -+ causal_global = causal_seqlens_kv_global[blk_coord[2]] - compute_common_params = SimpleNamespace( - blk_coord=blk_coord, - split_kv=split_kv, -@@ -1345,8 +1413,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - mAccO=mAccO, - mO=mO, - K=cache_seqs[blk_coord[2]], -+ causal_global=causal_global, -+ cp_rank=cp_rank, -+ split_start_key=k_index * self.mma_qk_tiler[1], - L=mCL.shape[1], -- H=mQL.shape[0], -+ H=self.get_valid_q_rows(blk_coord[1]), - tmem_ptr=tmem_ptr, - tidx=tidx, - tiled_mma_pv=tiled_mma_pv, -@@ -1366,11 +1437,74 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - p_cor_consumer_state=p_cor_consumer_state, - mma_o_consumer_state=mma_o_consumer_state, - ) -+ elif cutlass.const_expr(self.enable_dcp): -+ self.store_empty_dcp_work( -+ mO, -+ mLSE, -+ mAccO, -+ mAccLSE, -+ blk_coord, -+ self.get_valid_q_rows(blk_coord[1]), -+ tidx, -+ ) - tile_sched.advance_to_next_work() - work_tile = tile_sched.get_current_work() - - return - -+ @cute.jit -+ def store_empty_dcp_work( -+ self, -+ mO: Optional[cute.Tensor], -+ mLSE: Optional[cute.Tensor], -+ mAccO: Optional[cute.Tensor], -+ mAccLSE: Optional[cute.Tensor], -+ blk_coord: cute.Coord, -+ valid_q_rows: cutlass.Int32, -+ tidx: cutlass.Int32, -+ ): -+ """Materialize the neutral attention state for a split with no K tile.""" -+ cta_m_rows = self.mma_qk_tiler[0] // self.cluster_shape_mnk[0] -+ compute_threads = self.num_compute_warps * self.threads_per_warp -+ local_tidx = tidx % compute_threads -+ cta_row_base = blk_coord[0] * cta_m_rows -+ -+ for linear_idx in cutlass.range( -+ local_tidx, cta_m_rows * self.latent_dim, compute_threads -+ ): -+ cta_row = linear_idx // self.latent_dim -+ d_idx = linear_idx % self.latent_dim -+ q_row = cta_row_base + cta_row -+ if cute.elem_less(q_row, valid_q_rows): -+ if cutlass.const_expr(mAccO is None): -+ mO[q_row, d_idx, blk_coord[1], blk_coord[2]] = self.o_dtype(0.0) -+ else: -+ mAccO[ -+ q_row, -+ blk_coord[3], -+ d_idx, -+ blk_coord[1], -+ blk_coord[2], -+ ] = self.acc_dtype(0.0) -+ -+ if cute.elem_less(local_tidx, cta_m_rows): -+ q_row = cta_row_base + local_tidx -+ if cute.elem_less(q_row, valid_q_rows): -+ if cutlass.const_expr(mAccLSE is None): -+ mLSE[q_row, blk_coord[1], blk_coord[2]] = -self.lse_dtype.inf -+ else: -+ mAccLSE[ -+ q_row, blk_coord[3], blk_coord[1], blk_coord[2] -+ ] = -self.lse_dtype.inf -+ -+ @cute.jit -+ def get_valid_q_rows(self, q_tile_idx: cutlass.Int32) -> cutlass.Int32: -+ """Number of valid flattened query rows in one physical M tile.""" -+ valid_rows = self.mma_qk_tiler_mn[0] -+ if q_tile_idx == self.num_q_tiles - 1: -+ valid_rows = self.tail_q_rows -+ return valid_rows -+ - @cute.kernel - def reduction_kernel( - self, -@@ -1402,27 +1536,47 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - """ - bidx, bidy, bidz = cute.arch.block_idx() - tidx, _, _ = cute.arch.thread_idx() -- blk_coord = (bidx, bidy, bidz) -+ # Reducer blocks cover one D band of one real output row. Splitting -+ # D512 into independent D128 bands exposes enough CTAs for small-batch -+ # decode without any cross-CTA reduction or synchronization. -+ d_tile_idx = bidx % self.reducer_d_tiles -+ blk_coord = (bidx // self.reducer_d_tiles, bidy, bidz) -+ flat_q_row = blk_coord[1] * self.num_heads + blk_coord[0] -+ # The physical M tile is fixed at 128 rows, so map with shift/mask. -+ q_tile = flat_q_row >> 7 -+ q_tile_row = flat_q_row & 127 - local_split_kv = ( - block_split_kvs[blk_coord[2]] if self.is_var_split_kv else split_kv - ) - k_tile_total = cute.ceil_div(cache_seqs[blk_coord[2]], self.mma_qk_tiler[1]) -- k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) -- local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) -+ if cutlass.const_expr(self.enable_dcp): -+ k_tile_per_cta = cutlass.max( -+ cute.ceil_div(k_tile_total, local_split_kv), cutlass.Int32(1) -+ ) -+ local_split_kv = cutlass.max( -+ cute.ceil_div(k_tile_total, k_tile_per_cta), cutlass.Int32(1) -+ ) -+ else: -+ k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) -+ local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) - - # Alloc shared memory - smem = utils.SmemAllocator() -- storage = smem.allocate(MAX_SPLITS * self.acc_dtype.width // 8, 16) -+ storage = smem.allocate(self.reducer_max_splits * self.acc_dtype.width // 8, 16) - lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype) -- smem_lse_scale = cute.make_tensor(lse_scale_ptr, cute.make_layout(MAX_SPLITS)) -+ smem_lse_scale = cute.make_tensor( -+ lse_scale_ptr, cute.make_layout(self.reducer_max_splits) -+ ) - - if cutlass.const_expr(self.enable_pdl): - cute.arch.griddepcontrol_wait() -- gLSE = mAccLSE[blk_coord[0], None, blk_coord[1], blk_coord[2]] -+ gLSE = mAccLSE[q_tile_row, None, q_tile, blk_coord[2]] - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - if warp_idx == 0: - # calculate the global lse and exp ^ (local_lse - global_lse) -- lse_per_thread = cute.ceil_div(MAX_SPLITS, self.threads_per_warp) -+ lse_per_thread = cute.ceil_div( -+ self.reducer_max_splits, self.threads_per_warp -+ ) - - local_lse = cute.make_rmem_tensor( - cute.make_layout(lse_per_thread), self.lse_dtype -@@ -1439,39 +1593,61 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - # reduce the local lse - lse_max = cute.arch.fmax(lse_max, local_lse[i]) - lse_max = cute.arch.warp_reduction_max(lse_max) -- lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 -+ if cutlass.const_expr(self.enable_dcp): -+ has_valid_lse = lse_max != -self.lse_dtype.inf -+ lse_max = lse_max if has_valid_lse else 0.0 -+ else: -+ lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 - # calculate sum_lse - sum_lse = 0.0 - for i in cutlass.range_constexpr(lse_per_thread): - sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) - sum_lse = cute.arch.warp_reduction_sum(sum_lse) - # calculate the global_lse -- global_lse = ( -- lse_max + cute.math.log2(sum_lse, fastmath=True) -- if not sum_lse == self.lse_dtype(0.0) or sum_lse != sum_lse # noqa: SIM201 -- else self.lse_dtype.inf -- ) -- if tidx == 0: -- # Convert from kernel-internal log2 base to the natural-log -- # convention exposed to callers (matches trtllm-gen / flash-attn). -- # `1.0 / LOG2_E == ln(2)`. -- mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse * ( -- 1.0 / LOG2_E -+ if cutlass.const_expr(self.enable_dcp): -+ global_lse = ( -+ lse_max + cute.math.log2(sum_lse, fastmath=True) -+ if has_valid_lse -+ else -self.lse_dtype.inf - ) -+ else: -+ global_lse = ( -+ lse_max + cute.math.log2(sum_lse, fastmath=True) -+ if not sum_lse == self.lse_dtype(0.0) # noqa: SIM201 -+ or sum_lse != sum_lse -+ else self.lse_dtype.inf -+ ) -+ if d_tile_idx == 0: -+ if tidx == 0: -+ # Convert from kernel-internal log2 base to the natural-log -+ # convention exposed to callers (matches trtllm-gen / flash-attn). -+ # `1.0 / LOG2_E == ln(2)`. -+ mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse * ( -+ 1.0 / LOG2_E -+ ) - # store the scale to shared memory - for i in cutlass.range_constexpr(lse_per_thread): - split_kv_idx = tidx + i * self.threads_per_warp - if cute.elem_less(split_kv_idx, local_split_kv): -- smem_lse_scale[split_kv_idx] = cute.math.exp2( -- local_lse[i] - global_lse, fastmath=True -- ) -+ if cutlass.const_expr(self.enable_dcp): -+ smem_lse_scale[split_kv_idx] = ( -+ cute.math.exp2( -+ local_lse[i] - global_lse, fastmath=True -+ ) -+ if has_valid_lse -+ else self.acc_dtype(0.0) -+ ) -+ else: -+ smem_lse_scale[split_kv_idx] = cute.math.exp2( -+ local_lse[i] - global_lse, fastmath=True -+ ) - - pipeline.sync(barrier_id=4) - - elements_per_thread = cute.ceil_div( -- self.latent_dim, self.threads_per_warp * self.num_compute_warps -+ self.reducer_d_tile, self.threads_per_warp * self.num_compute_warps - ) -- gAccO = mAccO[blk_coord[0], None, None, blk_coord[1], blk_coord[2]] -+ gAccO = mAccO[q_tile_row, None, None, q_tile, blk_coord[2]] - rAccO = cute.make_rmem_tensor( - cute.make_layout(elements_per_thread), self.acc_dtype - ) -@@ -1479,11 +1655,19 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - rAccO.fill(0.0) - for i in range(local_split_kv): - for j in cutlass.range_constexpr(elements_per_thread): -- element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps -+ element_idx = ( -+ d_tile_idx * self.reducer_d_tile -+ + tidx -+ + j * self.threads_per_warp * self.num_compute_warps -+ ) - rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] - rO.store(rAccO.load().to(self.o_dtype)) - for j in cutlass.range_constexpr(elements_per_thread): -- element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps -+ element_idx = ( -+ d_tile_idx * self.reducer_d_tile -+ + tidx -+ + j * self.threads_per_warp * self.num_compute_warps -+ ) - mO[blk_coord[0], element_idx, blk_coord[1], blk_coord[2]] = rO[j] - if cutlass.const_expr(self.enable_pdl): - cute.arch.griddepcontrol_launch_dependents() -@@ -1665,7 +1849,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - # page table - mPT = common_params.mPT[None, common_params.blk_coord[2]] - -- # Flatten divide and partition global tensors for QK TMA load -+ # The single flattened descriptor is bounded by S_q * H. It can load -+ # across token boundaries and zero-fills only the final M128 tail. - # (bM, bK, rM, rK, rL) - mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2]) - gQL = cute.flat_divide(qk_params.mQL, mma_qk_tiler_mk) -@@ -1677,7 +1862,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - ) - tSgQL = thr_mma_qk.partition_A(gQL) - tSgQR = thr_mma_qk.partition_A(gQR) -- - cta_m = min( - qk_params.tiled_mma_qk.op.shape_mnk[0] - // qk_params.tiled_mma_qk.thr_id.shape, -@@ -1725,7 +1909,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - cute.group_modes(qk_params.sQ_rope, 0, 3), - cute.group_modes(tSgQR, 0, 3), - ) -- - tKCsKC, tCLgCL = cpasync.tma_partition( - qk_params.tma_atom_c_latent, - 0, -@@ -1742,11 +1925,14 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - tSgKR, - ) - -+ # Q is flattened as (Sq * Hq, D, B). The logical query-tile -+ # coordinate therefore lives in the residual-M mode created by -+ # flat_divide, rather than in a separate grouped-query mode. - tQLgQL = tQLgQL_mkl[ -- None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] -+ None, common_params.blk_coord[1], None, common_params.blk_coord[2] - ] - tQRgQR = tQRgQR_mkl[ -- None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] -+ None, common_params.blk_coord[1], None, common_params.blk_coord[2] - ] - - # Flatten divide and partition global tensors for V TMA load -@@ -1837,6 +2023,33 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - load_pt_release_state, - ) - -+ @cute.jit -+ def issue_q_tma_load( -+ self, -+ tma_atom_q_latent: cute.CopyAtom, -+ tQLgQL: cute.Tensor, -+ tQsQ: cute.Tensor, -+ tma_atom_q_rope: cute.CopyAtom, -+ tQRgQR: cute.Tensor, -+ tQsQ_rope: cute.Tensor, -+ tma_bar_ptr, -+ ): -+ """Issue one bounded flattened query-tile load into shared Q stages.""" -+ for i in cutlass.range(self.iterations_qk_latent): -+ cute.copy( -+ tma_atom_q_latent, -+ tQLgQL[None, i], -+ tQsQ[None, (i, 0)], -+ tma_bar_ptr=tma_bar_ptr, -+ ) -+ for i in cutlass.range(self.iterations_qk_rope): -+ cute.copy( -+ tma_atom_q_rope, -+ tQRgQR[None, i], -+ tQsQ_rope[None, i], -+ tma_bar_ptr=tma_bar_ptr, -+ ) -+ - @cute.jit - def load_tma_qk_one_k_tile( - self, -@@ -1893,22 +2106,15 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - tma_bar_ptr = common_params.load_q_pipeline.producer_get_barrier( - load_q_producer_state - ) -- for i in cutlass.range(self.iterations_qk_latent): -- # load q latent -- cute.copy( -- qk_params.tma_atom_q_latent, -- qk_params.tQLgQL[None, 0, i], -- qk_params.tQsQ[None, (i, 0)], -- tma_bar_ptr=tma_bar_ptr, -- ) -- for i in cutlass.range(self.iterations_qk_rope): -- # load q rope -- cute.copy( -- qk_params.tma_atom_q_rope, -- qk_params.tQRgQR[None, 0, i], -- qk_params.tQsQ_rope[None, i], -- tma_bar_ptr=tma_bar_ptr, -- ) -+ self.issue_q_tma_load( -+ qk_params.tma_atom_q_latent, -+ qk_params.tQLgQL, -+ qk_params.tQsQ, -+ qk_params.tma_atom_q_rope, -+ qk_params.tQRgQR, -+ qk_params.tQsQ_rope, -+ tma_bar_ptr, -+ ) - load_q_producer_state.advance() - load_kv_pipeline = common_params.load_kv_pipeline - tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) -@@ -2361,24 +2567,48 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] - """ - -- k_tile_total = cute.ceil_div(common_params.K, self.mma_qk_tiler[1]) -- - row_max = -self.acc_dtype.inf - row_sum = self.acc_dtype(0) - correction_factor = self.acc_dtype(1) - common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) - -- # Number of tiles from the global-K end that may contain causal-masked -- # positions. Min k_bound = K - (S_q-1), which can span up to -- # ceil((seq_len_q-2)/tile_N)+1 tiles (tile-boundary-crossing case). For -- # S_q=1 this reduces to 1 tile — identical to a plain K-bound check. -+ # The first tile that can contain a key at or beyond this query tile's -+ # earliest causal bound. A flat M tile starts at row q_tile * M, so -+ # floor(row / H) is its earliest query token. Later query tiles can -+ # therefore keep more fully dense K tiles on the compile-time unmasked -+ # path. The token division is outside the per-score loop (and is -+ # replicated by each participating CTA/compute group); tile_n is a -+ # static power of two, so the remaining division is a shift. - tile_n = self.mma_qk_tiler[1] -- mask_tile_count = (self.seq_len_q - 2 + tile_n - 1) // tile_n + 1 -- -- # first_mask_tile_idx is the global index of the first tile that may -- # need masking. Runtime because it depends on K (per-batch in -- # var-seq / split-KV). -- first_mask_tile_idx = k_tile_total - mask_tile_count -+ first_q_token = cutlass.Int32(0) -+ if cutlass.const_expr(self.num_q_tiles > 1): -+ first_q_token = ( -+ common_params.blk_coord[1] * self.mma_qk_tiler[0] -+ ) // self.num_heads -+ if cutlass.const_expr(self.enable_dcp): -+ # The earliest token in this full M128 query tile determines the -+ # dense-prefix boundary shared by both CTAs. Keep the physical -+ # local K bound separate so a partial final local tile is never -+ # misclassified as dense. -+ earliest_bound_numer = ( -+ common_params.causal_global -+ - common_params.cp_rank -+ - (self.seq_len_q - 1) -+ + first_q_token -+ ) -+ earliest_local_bound = cute.ceil_div( -+ cutlass.max(earliest_bound_numer, cutlass.Int32(0)), -+ self.cp_world, -+ ) -+ effective_local_bound = cutlass.min( -+ common_params.K, earliest_local_bound -+ ) -+ first_mask_tile_idx = effective_local_bound // tile_n -+ else: -+ first_mask_tile_idx = cutlass.max( -+ (common_params.K - self.seq_len_q + 1 + first_q_token) // tile_n, -+ cutlass.Int32(0), -+ ) - - # Phase 1: pure unmasked bulk tiles (all columns strictly < min k_bound). - while k_tile_count > 1 and k_index < first_mask_tile_idx: -@@ -2431,52 +2661,28 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - k_index = k_index + 1 - k_tile_count = k_tile_count - 1 - -- # Phase 3: this work-split's final tile. -- if cutlass.const_expr(common_params.mAccO is not None): -- # Split-KV: only apply mask when this final tile is globally in -- # the mask region (covers both last-split last-tile and straddling -- # splits). Runtime comparison. -- ( -- mma_s_consumer_state, -- p_mma_producer_state, -- p_cor_producer_state, -- row_max, -- row_sum, -- correction_factor, -- ) = self.softmax( -- common_params, -- softmax_params, -- k_index, -- mma_s_consumer_state, -- p_mma_producer_state, -- p_cor_producer_state, -- row_max, -- row_sum, -- correction_factor, -- k_index >= first_mask_tile_idx, -- True, -- ) -- else: -- ( -- mma_s_consumer_state, -- p_mma_producer_state, -- p_cor_producer_state, -- row_max, -- row_sum, -- correction_factor, -- ) = self.softmax( -- common_params, -- softmax_params, -- k_index, -- mma_s_consumer_state, -- p_mma_producer_state, -- p_cor_producer_state, -- row_max, -- row_sum, -- correction_factor, -- True, -- True, -- ) -+ # Phase 3: this work-split's final tile. Mask it only when its global -+ # index reaches the causal/K-bound region. -+ ( -+ mma_s_consumer_state, -+ p_mma_producer_state, -+ p_cor_producer_state, -+ row_max, -+ row_sum, -+ correction_factor, -+ ) = self.softmax( -+ common_params, -+ softmax_params, -+ k_index, -+ mma_s_consumer_state, -+ p_mma_producer_state, -+ p_cor_producer_state, -+ row_max, -+ row_sum, -+ correction_factor, -+ k_index >= first_mask_tile_idx, -+ True, -+ ) - - return mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state - -@@ -2635,8 +2841,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - :param correction_factor: The correction factor - :type correction_factor: cutlass.Float32 - :param apply_mask: Whether the tile needs K-bound / causal masking (Python bool -- for the unmasked/masked bulk loops; runtime cutlass.Boolean for the -- split-KV final iter where mask only applies on the global last tile). -+ for the unmasked/masked bulk loops; runtime cutlass.Boolean for a -+ work-split final tile at the exact per-query-tile mask boundary). - :type apply_mask: bool | cutlass.Boolean - :param is_local_last_tile: Whether the last tile is local - :type is_local_last_tile: cutlass.Boolean -@@ -2680,40 +2886,49 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype) - - row_max_new = row_max -- # Spec-decoding (MTP) causal mask: each row represents one (q_token, head) -- # pair; row r's effective K bound is K - (S_q - 1 - q_tok(r)). -- # With fold factor F = self.fold_sq_ratio (fold_sq=True), the M tile is -- # laid out as [F sub_q_tok][num_heads heads] and there are S_q/F outer -- # chunks indexed by blk_coord[1]: -- # q_tok(r) = blk_coord[1] * F + (r_global // num_heads) -- # r_global = row_in_cta + cluster_idx * (M_tile / cluster_m) -- # When fold_sq=False this reduces to q_tok = blk_coord[1]. For S_q=1 -- # this further reduces to k_bound = K (plain K-bound check). -- # Masked positions are filled with a large negative sentinel (not -inf) -- # to avoid NaN propagation when an entire row becomes masked. -+ # Spec-decoding (MTP) causal mask. A flattened row r corresponds to -+ # q_token=floor(r/H), whose last valid key is K-S_q+q_token. Avoid the -+ # per-element integer division using the equivalent integer predicate: -+ # -+ # H * (key_pos - K + S_q) <= r -+ # -+ # This arithmetic remains inside `if apply_mask`, so the compile-time -+ # dense path emits no row-mask work. Masked positions use a large -+ # negative sentinel (not -inf) to avoid all-masked-row NaNs. - cta_m_rows = self.mma_qk_tiler[0] // self.cluster_shape_mnk[0] - arch = BaseDSL._get_dsl().get_arch_enum() - if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f): - cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc) - for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): - if apply_mask: -- if cutlass.const_expr(self.fold_sq): -- q_tok = ( -- common_params.blk_coord[1] * self.fold_sq_ratio -- + (tTR_tS[i][0] + common_params.blk_coord[0] * cta_m_rows) -- // self.num_heads -+ flat_q_row = ( -+ common_params.blk_coord[1] * self.mma_qk_tiler[0] -+ + common_params.blk_coord[0] * cta_m_rows -+ + tTR_tS[i][0] -+ ) -+ key_pos = tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index -+ if cutlass.const_expr(self.enable_dcp): -+ mask_threshold = self.num_heads * ( -+ self.cp_world * key_pos -+ + common_params.cp_rank -+ - common_params.causal_global -+ + self.seq_len_q -+ ) -+ tTR_rAcc[i] = ( -+ tTR_rAcc[i] -+ if cute.elem_less(key_pos, common_params.K) -+ and not cute.elem_less(flat_q_row, mask_threshold) -+ else self.acc_dtype(-1.0e6) - ) - else: -- q_tok = common_params.blk_coord[1] -- k_bound = common_params.K - (self.seq_len_q - 1) + q_tok -- tTR_rAcc[i] = ( -- tTR_rAcc[i] -- if cute.elem_less( -- tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, -- k_bound, -+ mask_threshold = self.num_heads * ( -+ key_pos - common_params.K + self.seq_len_q -+ ) -+ tTR_rAcc[i] = ( -+ tTR_rAcc[i] -+ if not cute.elem_less(flat_q_row, mask_threshold) -+ else self.acc_dtype(-1.0e6) - ) -- else self.acc_dtype(-1.0e6) -- ) - # reduction for row_max - row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) - -@@ -2741,23 +2956,34 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout) - if apply_mask: - for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): -- if cutlass.const_expr(self.fold_sq): -- q_tok = ( -- common_params.blk_coord[1] * self.fold_sq_ratio -- + (tTR_tS[i][0] + common_params.blk_coord[0] * cta_m_rows) -- // self.num_heads -+ flat_q_row = ( -+ common_params.blk_coord[1] * self.mma_qk_tiler[0] -+ + common_params.blk_coord[0] * cta_m_rows -+ + tTR_tS[i][0] -+ ) -+ key_pos = tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index -+ if cutlass.const_expr(self.enable_dcp): -+ mask_threshold = self.num_heads * ( -+ self.cp_world * key_pos -+ + common_params.cp_rank -+ - common_params.causal_global -+ + self.seq_len_q -+ ) -+ tTR_rAcc[i] = ( -+ tTR_rAcc[i] -+ if cute.elem_less(key_pos, common_params.K) -+ and not cute.elem_less(flat_q_row, mask_threshold) -+ else self.acc_dtype(-1.0e6) - ) - else: -- q_tok = common_params.blk_coord[1] -- k_bound = common_params.K - (self.seq_len_q - 1) + q_tok -- tTR_rAcc[i] = ( -- tTR_rAcc[i] -- if cute.elem_less( -- tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, -- k_bound, -+ mask_threshold = self.num_heads * ( -+ key_pos - common_params.K + self.seq_len_q -+ ) -+ tTR_rAcc[i] = ( -+ tTR_rAcc[i] -+ if not cute.elem_less(flat_q_row, mask_threshold) -+ else self.acc_dtype(-1.0e6) - ) -- else self.acc_dtype(-1.0e6) -- ) - # reduction for row_max after manual masking - row_max_new = tTR_rAcc.load().reduce( - cute.ReductionOp.MAX, row_max_new, 0 -@@ -3173,15 +3399,45 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - # load o - cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) - -+ row_has_key = True -+ if cutlass.const_expr(self.enable_dcp): -+ # The coordinate tensor returned by local_tile already carries -+ # the CTA's offset within the full M128 query tile. -+ flat_q_row = ( -+ common_params.blk_coord[1] * self.mma_qk_tiler[0] -+ + tTR_cO[0][0] -+ ) -+ first_key_threshold = self.num_heads * ( -+ self.cp_world * common_params.split_start_key -+ + common_params.cp_rank -+ - common_params.causal_global -+ + self.seq_len_q -+ ) -+ row_has_key = cute.elem_less( -+ common_params.split_start_key, common_params.K -+ ) and not cute.elem_less(flat_q_row, first_key_threshold) -+ - # apply output scale and normalize by row_sum -- for i in cutlass.range( -- cute.size(tTR_rAcc), vectorize=True, unroll_full=True -- ): -- tTR_rAcc[i] = ( -- tTR_rAcc[i] -- * epilogue_params.output_scale -- * cute.arch.rcp_approx(row_sum) -+ if cutlass.const_expr(self.enable_dcp): -+ normalization_scale = ( -+ epilogue_params.output_scale * cute.arch.rcp_approx(row_sum) -+ ) -+ normalization_scale = ( -+ normalization_scale if row_has_key else self.acc_dtype(0.0) - ) -+ for i in cutlass.range( -+ cute.size(tTR_rAcc), vectorize=True, unroll_full=True -+ ): -+ tTR_rAcc[i] = tTR_rAcc[i] * normalization_scale -+ else: -+ for i in cutlass.range( -+ cute.size(tTR_rAcc), vectorize=True, unroll_full=True -+ ): -+ tTR_rAcc[i] = ( -+ tTR_rAcc[i] -+ * epilogue_params.output_scale -+ * cute.arch.rcp_approx(row_sum) -+ ) - - # store o to global memory - tR2G_rO_src = None -@@ -3262,6 +3518,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - cute.math.log2(row_sum, fastmath=True) - + epilogue_params.softmax_scale_log2 * row_max - ) -+ if cutlass.const_expr(self.enable_dcp): -+ lse = lse if row_has_key else -self.lse_dtype.inf - # When writing directly to the user-facing mLSE (single-tile, - # no split-KV merge), convert from log2 base to natural log. - # When writing the per-split intermediate (mAccLSE branch), keep -@@ -3504,25 +3762,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - - return tile_sched_params, grid - -- @staticmethod -- def compute_fold_sq_ratio(num_heads: int, seq_len_q: int, m_tile: int) -> int: -- """Derive the seq_len_q-into-heads fold factor F. -- -- Returns the largest integer F such that: -- - F divides seq_len_q evenly -- - num_heads * F ≤ m_tile -- - 1 ≤ F ≤ seq_len_q -- -- F=1 means no folding (i.e. ``fold_sq`` should be False at the caller). -- """ -- if num_heads >= m_tile: -- return 1 -- max_fold = min(seq_len_q, m_tile // num_heads) -- for f in range(max_fold, 0, -1): -- if seq_len_q % f == 0: -- return f -- return 1 -- - @staticmethod - def get_workspace_size( - H: int, -@@ -3552,12 +3791,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - """ - if split_kv == 1: - return 0 -- # Decode packs heads into a physical 128-wide MMA-M tile. For H < 128, -- # split-KV partials can still touch the padded head lanes before -- # reduction, so size the workspace for max(H, 128). Mirrors the same -- # padding applied in initialize_workspace(). See #3235. -- workspace_heads = max(H, 128) -- return B * workspace_heads * S * split_kv * (D + 1) * acc_dtype.width // 8 -+ # The first workspace mode is the physical MMA-M row extent. Direct -+ # callers may still pass a logical H < 128, so preserve the M128 floor -+ # used by initialize_workspace(). The flat internal path passes 128. -+ workspace_rows = max(H, 128) -+ return B * workspace_rows * S * split_kv * (D + 1) * acc_dtype.width // 8 - - @cute.jit - def initialize_workspace( -@@ -3593,29 +3831,29 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - """ - acc_o, acc_lse = None, None - if cutlass.const_expr(workspace is not None): -- # Pad head dim to the physical 128-wide MMA-M tile. Without this, -- # H<128 split-KV partials write past the workspace. See #3235. -- workspace_H = cutlass.max(H, cutlass.Int32(128)) -+ # The first workspace mode is the physical MMA-M row extent. -+ # Preserve the M128 floor for direct callers that pass H < 128. -+ workspace_rows = cutlass.max(H, cutlass.Int32(128)) - align = 256 // self.q_dtype.width - acc_o_layout = cute.make_layout( -- (workspace_H, split_kv, D, S, B), -+ (workspace_rows, split_kv, D, S, B), - stride=( - cute.assume(split_kv * D, align), - cute.assume(D, align), - 1, -- cute.assume(split_kv * workspace_H * D, align), -- cute.assume(workspace_H * split_kv * S * D, align), -+ cute.assume(split_kv * workspace_rows * D, align), -+ cute.assume(workspace_rows * split_kv * S * D, align), - ), - ) - acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=acc_dtype) - acc_o = cute.make_tensor(acc_o_iter, acc_o_layout) - acc_lse_layout = cute.make_layout( -- (workspace_H, split_kv, S, B), -+ (workspace_rows, split_kv, S, B), - stride=( - split_kv, - 1, -- workspace_H * split_kv, -- workspace_H * split_kv * S, -+ workspace_rows * split_kv, -+ workspace_rows * split_kv * S, - ), - ) - acc_lse_iter = cute.recast_ptr( -@@ -3697,11 +3935,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: - return False - if is_var_split_kv and not is_var_seq: - return False -- if mma_qk_tiler_mn[0] < H: -+ if H <= 0 or mma_qk_tiler_mn[0] < H: - return False -- # When H < M tile, fold up to F tokens of S into H (M_eff = H*F ≤ M_tile). -- # F is auto-picked as the largest divisor of S with H*F ≤ M_tile. -- # F=1 always works, so any (H ≤ M_tile, S ≥ 1) is implementable. -+ # Query (token, head) rows are packed continuously into 128-row M -+ # tiles. A partial final tile is supported, so any H <= M_tile and -+ # S >= 1 is valid. - if S < 1: - return False - if K <= 0: -diff --git a/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp8.py b/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp8.py -index f79683bc..f0aa6e72 100644 ---- a/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp8.py -+++ b/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp8.py -@@ -75,6 +75,7 @@ from cutlass.cutlass_dsl import BaseDSL - - from .mla_helpers import ( - ceil_div, -+ compute_q_tile_layout, - MAX_SPLITS, - LOG2_E, - MLAStaticTileScheduler, -@@ -106,15 +107,17 @@ launcher and ``flashinfer/cute_dsl/attention/mla_dispatch.py`` for impl selectio - - Constraints: - * Data type requirements: -- - Input/output: Float8E4M3FN -+ - Input: Float8E4M3FN -+ - Output: BFloat16 or Float8E4M3FN - - Accumulation and LSE: Float32 - * Fixed architecture parameters: -- - Number of attention heads: 128 -+ - Number of attention heads: 1-128 - - Latent dimension: 512 - - RoPE dimension: 64 - * Input query modes should be (NumHeads, LatentDim/RopeDim, SeqLenQ, BatchSize) - * Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize) --* Query sequence length must be 1-4 -+* Query sequence length must be positive; token/head rows are flattened across -+ 128-row M tiles with a safely padded final tile - * Only supports 2-CTA instructions - * Variable sequence length requires page table storage enabled - """ -@@ -136,7 +139,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - enable_pdl: bool, - num_heads: int = 128, - seq_len_q: int = 1, -- fold_sq: bool = False, -+ reducer_d_tiles: int = 1, -+ reducer_max_splits: int = MAX_SPLITS, -+ enable_dcp: bool = False, -+ cp_world: int = 1, - ): - """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel. - -@@ -162,17 +168,24 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - :type is_var_split_kv: bool - :param enable_pdl: Whether to use PDL - :type enable_pdl: bool -- :param num_heads: Number of attention heads (pre-fold). Used for the -- per-row spec-decoding (MTP) causal mask q_token_index computation. -+ :param num_heads: Number of attention heads. Defines the flattened -+ ``(query_token, head)`` row geometry and division-free causal mask. - :type num_heads: int -- :param seq_len_q: Query sequence length (pre-fold). Combined with -- ``num_heads`` to derive the per-row q_token used by the causal mask. -+ :param seq_len_q: Query sequence length. Combined with ``num_heads`` -+ to define the flattened query-row extent and causal boundary. - :type seq_len_q: int -- :param fold_sq: Whether to fold tokens of ``seq_len_q`` into the head -- dimension so the M tile becomes [F sub_q_tok][num_heads heads]. -- Required when ``num_heads < mma_qk_tiler_mn[0]`` and ``seq_len_q > 1`` -- so the M tile is fully populated. -- :type fold_sq: bool -+ :param reducer_d_tiles: Number of independent D bands reduced per row. -+ :type reducer_d_tiles: int -+ :param reducer_max_splits: Compile-time reducer capacity. Direct users -+ retain the generic 256-split capacity by default; callers choosing -+ a smaller specialization must cap runtime split-KV accordingly. -+ :type reducer_max_splits: int -+ :param enable_dcp: Statically enable decode context-parallel causal -+ masking over a cyclic rank-local KV shard. -+ :type enable_dcp: bool -+ :param cp_world: Compile-time DCP world size. Rank-local key ``k`` maps -+ to global key ``k * cp_world + cp_rank``. -+ :type cp_world: int - """ - - self.latent_dim = 512 -@@ -187,25 +200,36 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - self.page_size = page_size - self.is_var_seq = is_var_seq - self.is_var_split_kv = is_var_split_kv -+ if not 1 <= reducer_max_splits <= MAX_SPLITS: -+ raise ValueError( -+ f"reducer_max_splits must be in [1, {MAX_SPLITS}], " -+ f"got {reducer_max_splits}" -+ ) -+ if is_var_split_kv and reducer_max_splits != MAX_SPLITS: -+ raise ValueError( -+ "variable split-KV requires the generic reducer capacity " -+ f"of {MAX_SPLITS}, got {reducer_max_splits}" -+ ) -+ self.reducer_max_splits = reducer_max_splits - self.enable_pdl = enable_pdl -- # Original (pre-fold) num_heads and seq_len_q used for per-row -- # spec-decoding (MTP) causal q_token_index computation. When fold_sq is -- # True the M tile is laid out as [F sub_q_tok][num_heads heads]; the -- # full q_tok for row r is blk_coord[1] * F + (r // num_heads). -+ if cp_world < 1: -+ raise ValueError(f"cp_world must be positive, got {cp_world}") -+ if not enable_dcp and cp_world != 1: -+ raise ValueError( -+ "cp_world must be 1 when decode context parallelism is disabled" -+ ) -+ self.enable_dcp = enable_dcp -+ self.cp_world = cp_world -+ # Flatten query-token and head modes into one affine row space. An MMA -+ # M tile may cross token boundaries; only the final tile can be -+ # partially populated and is padded by the Q TMA OOB fill. - self.num_heads = num_heads - self.seq_len_q = seq_len_q -- # fold_sq (caller-controlled): whether the folding code path is enabled. -- # fold_sq_ratio (derived): fold factor F ≥ 1; the largest divisor of -- # seq_len_q with num_heads * F ≤ M_tile and F ≤ seq_len_q. When the -- # caller passes fold_sq=False, the kernel ignores the ratio. -- # When fold_sq=True but the derived ratio is 1, the folding branch -- # is taken with F=1 (a no-op transform). -- self.fold_sq = fold_sq -- self.fold_sq_ratio = ( -- BlackwellMultiHeadLatentAttentionForwardFP8.compute_fold_sq_ratio( -- num_heads, seq_len_q, mma_qk_tiler_mn[0] -- ) -- ) -+ ( -+ self.total_q_rows, -+ self.num_q_tiles, -+ self.tail_q_rows, -+ ) = compute_q_tile_layout(num_heads, seq_len_q, mma_qk_tiler_mn[0]) - self.cluster_shape_mnk = (2, 1, 1) - self.use_2cta_instrs = True - # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2), -@@ -213,6 +237,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - self.warps_in_n = 2 - self.num_compute_warps = 4 - self.threads_per_warp = 32 -+ if reducer_d_tiles not in (1, 2, 4): -+ raise ValueError(f"unsupported reducer_d_tiles={reducer_d_tiles}") -+ self.reducer_d_tiles = reducer_d_tiles -+ self.reducer_d_tile = self.latent_dim // reducer_d_tiles - mma_qk_tiler_k = self.rope_dim * 2 - self.mma_qk_tiler = ( - self.mma_qk_tiler_mn[0], -@@ -294,13 +322,13 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - num_threads=(self.threads_per_warp * self.num_compute_warps), - ) - # Pingpong order barriers (OrderedSequenceBarrier<1,2> pattern). Each -- # group waits on its own bar via arrive_and_wait (= bar.sync); the OTHER -- # group signals via split-phase .arrive() (non-blocking). num_threads -+ # group waits on its own bar via arrive_and_wait (= bar.sync), while the -+ # other group normally signals via split-phase .arrive(). num_threads - # MUST cover both groups (256 = wait-side 128 + signal-side 128) so the - # bar releases only after BOTH have arrived — that's the cross-group -- # serialization that gives the TMEM peer-read its happens-before. -- # Init-phase trick: g1 pre-arrives bar_0 once at warp setup so g0's -- # first arrive_and_wait completes without waiting for g1's loop arrive. -+ # serialization that gives the TMEM peer-read its happens-before. The -+ # final owner also waits, closing the named-barrier generation before a -+ # persistent group can start its next logical work item. - self.softmax_order_bar_0 = pipeline.NamedBarrier( - barrier_id=5, - num_threads=(self.threads_per_warp * self.num_total_compute_warps), -@@ -353,6 +381,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - workspace: cute.Tensor, - split_kv: cutlass.Int32, - cache_seqs: Optional[cute.Tensor], -+ causal_seqlens_kv_global: Optional[cute.Tensor], -+ cp_rank: cutlass.Int32, - block_split_kvs: Optional[cute.Tensor], - softmax_scale: cutlass.Float32, - output_scale: cutlass.Float32, -@@ -388,6 +418,12 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - :type split_kv: cutlass.Int32 - :param cache_seqs: The cache sequences tensor with shape [batch_size] - :type cache_seqs: cute.Tensor -+ :param causal_seqlens_kv_global: Global exclusive causal bounds for the -+ newest query, with shape [batch_size]. Used only when DCP is -+ statically enabled. -+ :type causal_seqlens_kv_global: cute.Tensor -+ :param cp_rank: Runtime DCP rank in ``[0, cp_world)``. -+ :type cp_rank: cutlass.Int32 - :param block_split_kvs: The block split KV tensor with shape [batch_size] - :type block_split_kvs: cute.Tensor - :param softmax_scale: The scale factor for softmax -@@ -428,6 +464,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - q_latent = _reinterpret_4d(q_latent) - q_rope = _reinterpret_4d(q_rope) - o = _reinterpret_4d(o) -+ # Keep unpacked output aliases for the split-KV reducer, which -+ # launches only real (head, q_token) rows. -+ o_unpacked = o - - # Reinterpret contiguous [num_pages, page_size, D] as [page_size, D, num_pages] - # Input stride: (PS*D, D, 1) → Target: (D, 1, PS*D) -@@ -461,54 +500,53 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - stride=(lse.stride[2], lse.stride[1], lse.stride[0]), - ), - ) -+ lse_unpacked = lse - -- # When num_heads < M tile, fold up to F = fold_sq_ratio tokens of -- # seq_len_q into the head dimension so M_eff = num_heads * F (≤ M_tile). -- # E.g., H=32, S_q=4 → F=4, M_eff=128, S_q_eff=1 -- # E.g., H=32, S_q=8 → F=4, M_eff=128, S_q_eff=2 -- # This works because MLA shares KV across all heads/queries independently. -- # Tensor layout: [H, D, S_q, B] → [H*F, D, S_q/F, B]; relies on -- # stride_S == stride_H * H (always true for contiguous [B, S_q, H, D] -- # tensors after _reinterpret_4d). -- if cutlass.const_expr(self.fold_sq): -- F = self.fold_sq_ratio -- -- def _fold_sq_4d(t): -- return cute.make_tensor( -- t.iterator, -- cute.make_layout( -- ( -- t.shape[0] * F, -- t.shape[1], -- t.shape[2] // F, -- t.shape[3], -- ), -- stride=( -- t.stride[0], -- t.stride[1], -- t.stride[2] * F, -- t.stride[3], -- ), -- ), -- ) -- -- q_latent = _fold_sq_4d(q_latent) -- q_rope = _fold_sq_4d(q_rope) -- o = _fold_sq_4d(o) -- # LSE: [H, S_q, B] → [H*F, S_q/F, B] -- lse = cute.make_tensor( -- lse.iterator, -+ # Flatten [H, D, S_q, B] into the affine row order -+ # flat_row = q_token * H + q_head. A single bounded TMA descriptor can -+ # then cross token boundaries and OOB-zero-fill only the final M tile. -+ def _flatten_q_rows(t): -+ return cute.make_tensor( -+ t.iterator, - cute.make_layout( -- (lse.shape[0] * F, lse.shape[1] // F, lse.shape[2]), -- stride=(lse.stride[0], lse.stride[1] * F, lse.stride[2]), -+ (self.total_q_rows, t.shape[1], t.shape[3]), -+ stride=(t.stride[0], t.stride[1], t.stride[3]), - ), - ) - -+ q_latent = _flatten_q_rows(q_latent) -+ q_rope = _flatten_q_rows(q_rope) -+ -+ # O/LSE expose one physical M tile per query tile. Their final tile's -+ # virtual padding is never dereferenced because stores are -+ # predicated by get_valid_q_rows(). -+ m_tile = self.mma_qk_tiler_mn[0] -+ runtime_num_q_tiles = cute.ceil_div(o.shape[0] * o.shape[2], m_tile) -+ o = cute.make_tensor( -+ o.iterator, -+ cute.make_layout( -+ (m_tile, o.shape[1], runtime_num_q_tiles, o.shape[3]), -+ stride=( -+ o.stride[0], -+ o.stride[1], -+ o.stride[0] * m_tile, -+ o.stride[3], -+ ), -+ ), -+ ) -+ lse = cute.make_tensor( -+ lse.iterator, -+ cute.make_layout( -+ (m_tile, runtime_num_q_tiles, lse.shape[2]), -+ stride=(lse.stride[0], lse.stride[0] * m_tile, lse.stride[2]), -+ ), -+ ) -+ - acc_o, acc_lse = self.initialize_workspace( -- q_latent.shape[0], -- q_latent.shape[1], -- q_latent.shape[2], -- q_latent.shape[3], -+ o.shape[0], -+ o.shape[1], -+ o.shape[2], -+ o.shape[3], - split_kv, - self.acc_dtype, - workspace, -@@ -841,6 +879,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - acc_lse, - split_kv, - cache_seqs, -+ causal_seqlens_kv_global, -+ cp_rank, - block_split_kvs, - softmax_scale_log2, - output_scale, -@@ -867,17 +907,21 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - ) - if cutlass.const_expr(acc_o is not None): - self.reduction_kernel( -- o, -- lse, -+ o_unpacked, -+ lse_unpacked, - acc_o, - acc_lse, - split_kv, - cache_seqs, - block_split_kvs, - ).launch( -- grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]), -+ grid=( -+ o_unpacked.shape[0] * self.reducer_d_tiles, -+ o_unpacked.shape[2], -+ o_unpacked.shape[3], -+ ), - block=[self.threads_per_warp * self.num_compute_warps, 1, 1], -- smem=MAX_SPLITS * self.acc_dtype.width // 8, -+ smem=self.reducer_max_splits * self.acc_dtype.width // 8, - stream=stream, - min_blocks_per_mp=1, - use_pdl=self.enable_pdl, -@@ -943,6 +987,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - mAccLSE: Optional[cute.Tensor], - split_kv: cutlass.Int32, - cache_seqs: cute.Tensor, -+ causal_seqlens_kv_global: Optional[cute.Tensor], -+ cp_rank: cutlass.Int32, - block_split_kvs: cute.Tensor, - softmax_scale_log2: cutlass.Float32, - output_scale: cutlass.Float32, -@@ -985,8 +1031,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - :type mQL: cute.Tensor - :param tma_atom_q_rope: TMA copy atom for query rope tensor - :type tma_atom_q_rope: cute.CopyAtom -- :param mKR: Compressed rope tensor -- :type mKR: cute.Tensor -+ :param mQR: query rope tensor -+ :type mQR: cute.Tensor - :param tma_atom_c_latent: TMA copy atom for c latent tensor - :type tma_atom_c_latent: cute.CopyAtom - :param mCL: Compressed latent tensor -@@ -1009,6 +1055,12 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - :type split_kv: cutlass.Int32 - :param cache_seqs: The variable sequence length tensor - :type cache_seqs: cute.Tensor -+ :param causal_seqlens_kv_global: Global exclusive causal bounds for the -+ newest query. Present only in the statically enabled DCP -+ specialization. -+ :type causal_seqlens_kv_global: cute.Tensor -+ :param cp_rank: Runtime rank of this cyclic KV shard -+ :type cp_rank: cutlass.Int32 - :param block_split_kvs: The per-block split_kv values tensor - :type block_split_kvs: cute.Tensor - :param softmax_scale_log2: The log2 scale factor for softmax -@@ -1425,6 +1477,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - split_kv, cache_seqs, block_split_kvs, blk_coord - ) - if k_tile_count > 0: -+ causal_seq_len = cutlass.Int32(0) -+ if cutlass.const_expr(self.enable_dcp): -+ causal_seq_len = causal_seqlens_kv_global[blk_coord[2]] - compute_common_params = SimpleNamespace( - blk_coord=blk_coord, - split_kv=split_kv, -@@ -1433,6 +1488,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - mAccO=mAccO, - mO=mO, - K=cache_seqs[blk_coord[2]], -+ causal_seq_len=causal_seq_len, -+ cp_rank=cp_rank, - L=mCL.shape[1], - tmem_ptr=tmem_ptr, - tidx=tidx, -@@ -1492,17 +1549,18 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - mma_s_consumer_state.advance() - p_mma_producer_state.advance() - p_cor_producer_state.advance() -- # Pingpong init-phase trick: g1 pre-arrives softmax_order_bar_0 once -- # so g0's FIRST softmax_order_bar_0.arrive_and_wait() completes -- # without waiting for g1's first loop arrive. Replaces a per-iter -- # is_first_tile guard. Bar is num_total_compute_warps (256 threads) -- # so wait-side 128 + signal-side 128 = 256 → release. -+ # Each nonempty compute() call initializes g1's peer metadata and -+ # pre-arrives softmax_order_bar_0 for that logical work item. This -+ # lets g0's first wait complete without a first-tile special case. - while work_tile.is_valid_tile: - blk_coord = work_tile.tile_idx - k_index, k_tile_count, local_split_kv = self.get_k_tile_count( - split_kv, cache_seqs, block_split_kvs, blk_coord - ) - if k_tile_count > 0: -+ causal_seq_len = cutlass.Int32(0) -+ if cutlass.const_expr(self.enable_dcp): -+ causal_seq_len = causal_seqlens_kv_global[blk_coord[2]] - compute_common_params = SimpleNamespace( - blk_coord=blk_coord, - split_kv=split_kv, -@@ -1511,6 +1569,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - mAccO=mAccO, - mO=mO, - K=cache_seqs[blk_coord[2]], -+ causal_seq_len=causal_seq_len, -+ cp_rank=cp_rank, - L=mCL.shape[1], - tmem_ptr=tmem_ptr, - tidx=tidx, -@@ -1568,16 +1628,22 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - split_kv, cache_seqs, block_split_kvs, blk_coord - ) - if k_tile_count > 0: -+ causal_seq_len = cutlass.Int32(0) -+ if cutlass.const_expr(self.enable_dcp): -+ causal_seq_len = causal_seqlens_kv_global[blk_coord[2]] - compute_common_params = SimpleNamespace( - blk_coord=blk_coord, - split_kv=split_kv, - local_split_kv=local_split_kv, -+ k_index=k_index, - smem_exchange=epilogue_smem_exchange, - mAccO=mAccO, - mO=mO, - K=cache_seqs[blk_coord[2]], -+ causal_seq_len=causal_seq_len, -+ cp_rank=cp_rank, - L=mCL.shape[1], -- H=mQL.shape[0], -+ H=self.get_valid_q_rows(blk_coord[1]), - tmem_ptr=tmem_ptr, - tidx=tidx, - tiled_mma_pv=tiled_mma_pv, -@@ -1597,10 +1663,73 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - p_cor_consumer_state=p_cor_consumer_state, - mma_o_consumer_state=mma_o_consumer_state, - ) -+ elif cutlass.const_expr(self.enable_dcp): -+ self.store_empty_dcp_work( -+ mO, -+ mLSE, -+ mAccO, -+ mAccLSE, -+ blk_coord, -+ self.get_valid_q_rows(blk_coord[1]), -+ tidx, -+ ) - tile_sched.advance_to_next_work() - work_tile = tile_sched.get_current_work() - return - -+ @cute.jit -+ def store_empty_dcp_work( -+ self, -+ mO: Optional[cute.Tensor], -+ mLSE: Optional[cute.Tensor], -+ mAccO: Optional[cute.Tensor], -+ mAccLSE: Optional[cute.Tensor], -+ blk_coord: cute.Coord, -+ valid_q_rows: cutlass.Int32, -+ tidx: cutlass.Int32, -+ ): -+ """Materialize O=0 and LSE=-inf for a DCP split with no physical K tile.""" -+ cta_m_rows = self.mma_qk_tiler[0] // self.cluster_shape_mnk[0] -+ compute_threads = self.num_compute_warps * self.threads_per_warp -+ local_tidx = tidx % compute_threads -+ cta_row_base = blk_coord[0] * cta_m_rows -+ -+ for linear_idx in cutlass.range( -+ local_tidx, cta_m_rows * self.latent_dim, compute_threads -+ ): -+ cta_row = linear_idx // self.latent_dim -+ d_idx = linear_idx % self.latent_dim -+ q_row = cta_row_base + cta_row -+ if cute.elem_less(q_row, valid_q_rows): -+ if cutlass.const_expr(mAccO is None): -+ mO[q_row, d_idx, blk_coord[1], blk_coord[2]] = self.o_dtype(0.0) -+ else: -+ mAccO[ -+ q_row, -+ blk_coord[3], -+ d_idx, -+ blk_coord[1], -+ blk_coord[2], -+ ] = self.acc_dtype(0.0) -+ -+ if cute.elem_less(local_tidx, cta_m_rows): -+ q_row = cta_row_base + local_tidx -+ if cute.elem_less(q_row, valid_q_rows): -+ if cutlass.const_expr(mAccLSE is None): -+ mLSE[q_row, blk_coord[1], blk_coord[2]] = -self.lse_dtype.inf -+ else: -+ mAccLSE[ -+ q_row, blk_coord[3], blk_coord[1], blk_coord[2] -+ ] = -self.lse_dtype.inf -+ -+ @cute.jit -+ def get_valid_q_rows(self, q_tile_idx: cutlass.Int32) -> cutlass.Int32: -+ """Number of valid flattened query rows in one physical M tile.""" -+ valid_rows = self.mma_qk_tiler_mn[0] -+ if q_tile_idx == self.num_q_tiles - 1: -+ valid_rows = self.tail_q_rows -+ return valid_rows -+ - @cute.kernel - def reduction_kernel( - self, -@@ -1632,27 +1761,47 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - """ - bidx, bidy, bidz = cute.arch.block_idx() - tidx, _, _ = cute.arch.thread_idx() -- blk_coord = (bidx, bidy, bidz) -+ # Reducer blocks cover one D band of one real output row. Splitting -+ # D512 into independent D128 bands exposes enough CTAs for small-batch -+ # decode without any cross-CTA reduction or synchronization. -+ d_tile_idx = bidx % self.reducer_d_tiles -+ blk_coord = (bidx // self.reducer_d_tiles, bidy, bidz) -+ flat_q_row = blk_coord[1] * self.num_heads + blk_coord[0] -+ # The physical M tile is fixed at 128 rows, so map with shift/mask. -+ q_tile = flat_q_row >> 7 -+ q_tile_row = flat_q_row & 127 - local_split_kv = ( - block_split_kvs[blk_coord[2]] if self.is_var_split_kv else split_kv - ) - k_tile_total = cute.ceil_div(cache_seqs[blk_coord[2]], self.mma_qk_tiler[1]) -- k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) -- local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) -+ if cutlass.const_expr(self.enable_dcp): -+ k_tile_per_cta = cutlass.max( -+ cute.ceil_div(k_tile_total, local_split_kv), cutlass.Int32(1) -+ ) -+ local_split_kv = cutlass.max( -+ cute.ceil_div(k_tile_total, k_tile_per_cta), cutlass.Int32(1) -+ ) -+ else: -+ k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) -+ local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) - - # Alloc shared memory - smem = utils.SmemAllocator() -- storage = smem.allocate(MAX_SPLITS * self.acc_dtype.width // 8, 16) -+ storage = smem.allocate(self.reducer_max_splits * self.acc_dtype.width // 8, 16) - lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype) -- smem_lse_scale = cute.make_tensor(lse_scale_ptr, cute.make_layout(MAX_SPLITS)) -+ smem_lse_scale = cute.make_tensor( -+ lse_scale_ptr, cute.make_layout(self.reducer_max_splits) -+ ) - - if cutlass.const_expr(self.enable_pdl): - cute.arch.griddepcontrol_wait() -- gLSE = mAccLSE[blk_coord[0], None, blk_coord[1], blk_coord[2]] -+ gLSE = mAccLSE[q_tile_row, None, q_tile, blk_coord[2]] - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - if warp_idx == 0: - # calculate the global lse and exp ^ (local_lse - global_lse) -- lse_per_thread = cute.ceil_div(MAX_SPLITS, self.threads_per_warp) -+ lse_per_thread = cute.ceil_div( -+ self.reducer_max_splits, self.threads_per_warp -+ ) - - local_lse = cute.make_rmem_tensor( - cute.make_layout(lse_per_thread), self.lse_dtype -@@ -1669,39 +1818,59 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - # reduce the local lse - lse_max = cute.arch.fmax(lse_max, local_lse[i]) - lse_max = cute.arch.warp_reduction_max(lse_max) -- lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 -+ if cutlass.const_expr(self.enable_dcp): -+ has_valid_lse = lse_max != -self.lse_dtype.inf -+ lse_max = lse_max if has_valid_lse else 0.0 -+ else: -+ lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 - # calculate sum_lse - sum_lse = 0.0 - for i in cutlass.range_constexpr(lse_per_thread): - sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) - sum_lse = cute.arch.warp_reduction_sum(sum_lse) - # calculate the global_lse -- global_lse = ( -- lse_max + cute.math.log2(sum_lse, fastmath=True) -- if not sum_lse == self.lse_dtype(0.0) or sum_lse != sum_lse # noqa: SIM201 -- else self.lse_dtype.inf -- ) -- if tidx == 0: -- # Convert from kernel-internal log2 base to the natural-log -- # convention exposed to callers (matches trtllm-gen / flash-attn). -- # `1.0 / LOG2_E == ln(2)`. -- mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse * ( -- 1.0 / LOG2_E -+ if cutlass.const_expr(self.enable_dcp): -+ global_lse = ( -+ lse_max + cute.math.log2(sum_lse, fastmath=True) -+ if has_valid_lse -+ else -self.lse_dtype.inf -+ ) -+ else: -+ global_lse = ( -+ lse_max + cute.math.log2(sum_lse, fastmath=True) -+ if not sum_lse == self.lse_dtype(0.0) # noqa: SIM201 -+ or sum_lse != sum_lse -+ else self.lse_dtype.inf - ) -+ if d_tile_idx == 0: -+ if tidx == 0: -+ # Convert from kernel-internal log2 base to the natural-log -+ # convention exposed to callers (matches trtllm-gen / flash-attn). -+ # `1.0 / LOG2_E == ln(2)`. -+ mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse * ( -+ 1.0 / LOG2_E -+ ) - # store the scale to shared memory - for i in cutlass.range_constexpr(lse_per_thread): - split_kv_idx = tidx + i * self.threads_per_warp - if cute.elem_less(split_kv_idx, local_split_kv): -- smem_lse_scale[split_kv_idx] = cute.math.exp2( -- local_lse[i] - global_lse, fastmath=True -- ) -+ if cutlass.const_expr(self.enable_dcp): -+ smem_lse_scale[split_kv_idx] = ( -+ cute.math.exp2(local_lse[i] - global_lse, fastmath=True) -+ if has_valid_lse -+ else self.acc_dtype(0.0) -+ ) -+ else: -+ smem_lse_scale[split_kv_idx] = cute.math.exp2( -+ local_lse[i] - global_lse, fastmath=True -+ ) - - pipeline.sync(barrier_id=4) - - elements_per_thread = cute.ceil_div( -- self.latent_dim, self.threads_per_warp * self.num_compute_warps -+ self.reducer_d_tile, self.threads_per_warp * self.num_compute_warps - ) -- gAccO = mAccO[blk_coord[0], None, None, blk_coord[1], blk_coord[2]] -+ gAccO = mAccO[q_tile_row, None, None, q_tile, blk_coord[2]] - rAccO = cute.make_rmem_tensor( - cute.make_layout(elements_per_thread), self.acc_dtype - ) -@@ -1709,11 +1878,19 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - rAccO.fill(0.0) - for i in range(local_split_kv): - for j in cutlass.range_constexpr(elements_per_thread): -- element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps -+ element_idx = ( -+ d_tile_idx * self.reducer_d_tile -+ + tidx -+ + j * self.threads_per_warp * self.num_compute_warps -+ ) - rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] - rO.store(rAccO.load().to(self.o_dtype)) - for j in cutlass.range_constexpr(elements_per_thread): -- element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps -+ element_idx = ( -+ d_tile_idx * self.reducer_d_tile -+ + tidx -+ + j * self.threads_per_warp * self.num_compute_warps -+ ) - mO[blk_coord[0], element_idx, blk_coord[1], blk_coord[2]] = rO[j] - if cutlass.const_expr(self.enable_pdl): - cute.arch.griddepcontrol_launch_dependents() -@@ -1820,7 +1997,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - # page table - mPT = common_params.mPT[None, common_params.blk_coord[2]] - -- # Flatten divide and partition global tensors for QK TMA load -+ # The descriptor has one bounded flat M mode, so TMA can cross query -+ # token boundaries and zero-fills only the final physical tile. - # (bM, bK, rM, rK, rL) - mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2]) - gQL = cute.flat_divide(qk_params.mQL, mma_qk_tiler_mk) -@@ -1898,11 +2076,14 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - tSgKR, - ) - -+ # Q is flattened as (Sq * Hq, D, B). The logical query-tile -+ # coordinate therefore lives in the residual-M mode created by -+ # flat_divide, rather than in a separate grouped-query mode. - tQLgQL = tQLgQL_mkl[ -- None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] -+ None, common_params.blk_coord[1], None, common_params.blk_coord[2] - ] - tQRgQR = tQRgQR_mkl[ -- None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] -+ None, common_params.blk_coord[1], None, common_params.blk_coord[2] - ] - - # set extra params -@@ -2002,6 +2183,33 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - k_tile_count -= 1 - return load_v_producer_state - -+ @cute.jit -+ def issue_q_tma_load( -+ self, -+ tma_atom_q_latent: cute.CopyAtom, -+ tQLgQL: cute.Tensor, -+ tQsQ: cute.Tensor, -+ tma_atom_q_rope: cute.CopyAtom, -+ tQRgQR: cute.Tensor, -+ tQsQ_rope: cute.Tensor, -+ tma_bar_ptr, -+ ): -+ """Issue one bounded flattened query-tile load into shared Q stages.""" -+ for i in cutlass.range_constexpr(self.iterations_qk_latent): -+ cute.copy( -+ tma_atom_q_latent, -+ tQLgQL[None, i], -+ tQsQ[None, (i, 0)], -+ tma_bar_ptr=tma_bar_ptr, -+ ) -+ for i in cutlass.range_constexpr(self.iterations_qk_rope): -+ cute.copy( -+ tma_atom_q_rope, -+ tQRgQR[None, i], -+ tQsQ_rope[None, i], -+ tma_bar_ptr=tma_bar_ptr, -+ ) -+ - @cute.jit - def load_tma_qk_one_k_tile( - self, -@@ -2057,22 +2265,15 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - tma_bar_ptr = load_q_pipeline.producer_get_barrier(load_q_producer_state) - # expect the extra bytes for q. - load_q_pipeline.producer_acquire(load_q_producer_state) -- for i in cutlass.range_constexpr(self.iterations_qk_latent): -- # load q latent -- cute.copy( -- qk_params.tma_atom_q_latent, -- qk_params.tQLgQL[None, 0, i], -- qk_params.tQsQ[None, (i, 0)], -- tma_bar_ptr=tma_bar_ptr, -- ) -- for i in cutlass.range_constexpr(self.iterations_qk_rope): -- # load q rope -- cute.copy( -- qk_params.tma_atom_q_rope, -- qk_params.tQRgQR[None, 0, i], -- qk_params.tQsQ_rope[None, i], -- tma_bar_ptr=tma_bar_ptr, -- ) -+ self.issue_q_tma_load( -+ qk_params.tma_atom_q_latent, -+ qk_params.tQLgQL, -+ qk_params.tQsQ, -+ qk_params.tma_atom_q_rope, -+ qk_params.tQRgQR, -+ qk_params.tQsQ_rope, -+ tma_bar_ptr, -+ ) - load_q_producer_state.advance() - # get the mbar ptr from pipeline. - tma_bar_ptr = common_params.load_k_pipeline.producer_get_barrier( -@@ -2525,7 +2726,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - row_max = self.acc_dtype(self.init_row_max) - row_sum = self.acc_dtype(0) - correction_factor = self.acc_dtype(1) -- odd_k_tile = k_tile_count % 2 == 1 -+ k_tile_count_init = k_tile_count -+ odd_k_tile = k_tile_count_init % 2 == 1 - # 2softmax: g0 takes even k-tiles, g1 takes odd k-tiles. g1 advances all - # pipeline states once at entry to start on stage 1 (one stage per group). - if cutlass.const_expr(is_second_compute_warp): -@@ -2544,26 +2746,59 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - common_params, softmax_params, p_cor_producer_state - ) - self.softmax_order_bar_0.arrive() -- # Number of tiles from the global-K end that may contain causal-masked -- # positions. Min k_bound = K - (S_q-1), which can span up to -- # ceil((seq_len_q-2)/tile_N)+1 tiles (tile-boundary-crossing case). For -- # S_q=1 this reduces to 1 tile — identical to a plain K-bound check. -+ # The first tile that can contain a key at or beyond this query tile's -+ # earliest causal bound. A flat M tile starts at row q_tile * M, so -+ # floor(row / H) is its earliest query token. Later query tiles can -+ # therefore keep more fully dense K tiles on the compile-time unmasked -+ # path. The token division is outside the per-score loop (and is -+ # replicated by each participating CTA/compute group); tile_n is a -+ # static power of two, so the remaining division is a shift. - tile_n = self.mma_qk_tiler[1] -- mask_tile_count = (self.seq_len_q - 2 + tile_n - 1) // tile_n + 1 -+ first_q_token = cutlass.Int32(0) -+ if cutlass.const_expr(self.num_q_tiles > 1): -+ first_q_token = ( -+ common_params.blk_coord[1] * self.mma_qk_tiler[0] -+ ) // self.num_heads -+ if cutlass.const_expr(self.enable_dcp): -+ # The earliest chronological query represented by this physical -+ # M128 tile has the smallest visible rank-local K bound. Compute -+ # that bound once per query tile, then cap it by the physical -+ # rank-local K extent so a partial physical tail is still masked. -+ dcp_bound_numerator = cutlass.max( -+ common_params.causal_seq_len -+ - common_params.cp_rank -+ - (self.seq_len_q - 1) -+ + first_q_token, -+ cutlass.Int32(0), -+ ) -+ earliest_local_bound = ( -+ dcp_bound_numerator + self.cp_world - 1 -+ ) // self.cp_world -+ effective_local_bound = cutlass.min(common_params.K, earliest_local_bound) -+ first_mask_tile_idx = effective_local_bound // tile_n -+ else: -+ first_mask_tile_idx = cutlass.max( -+ (common_params.K - self.seq_len_q + 1 + first_q_token) // tile_n, -+ cutlass.Int32(0), -+ ) - -- # first_mask_tile_idx is the global index of the first tile that may -- # need masking. Runtime because it depends on K (per-batch in -- # var-seq / split-KV). -- first_mask_tile_idx = k_tile_total - mask_tile_count -+ # The non-split two-softmax path still needs its global final tile in -+ # phase 2 for final metadata exchange, even when that tile is dense. -+ # Masking itself remains controlled by first_mask_tile_idx below. -+ first_phase_2_tile_idx = cutlass.min(first_mask_tile_idx, k_tile_total - 1) - - # Phase 1: pure unmasked bulk tiles (all columns strictly < min k_bound). - # 2softmax: each group steps by 2 k-tiles; phase boundaries respect this. -- while k_tile_count > 0 and k_index < first_mask_tile_idx: -+ while k_tile_count > 0 and k_index < first_phase_2_tile_idx: - is_local_last_tile = ( - False - if cutlass.const_expr(common_params.mAccO is None) - else k_tile_count == 1 - ) -+ if cutlass.const_expr(is_second_compute_warp): -+ is_order_chain_last_tile = k_tile_count == 1 and not odd_k_tile -+ else: -+ is_order_chain_last_tile = k_tile_count == 1 and odd_k_tile - ( - mma_s_consumer_state, - p_mma_producer_state, -@@ -2584,6 +2819,7 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - is_second_compute_warp, - False, - is_local_last_tile, -+ is_order_chain_last_tile, - ) - k_index = k_index + 2 - k_tile_count = k_tile_count - 1 -@@ -2600,6 +2836,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - # Phase 2: remaining tiles that overlap the causal / K-bound region, - # including this work-split's final tile. - while k_tile_count > 0: -+ if cutlass.const_expr(is_second_compute_warp): -+ is_order_chain_last_tile = k_tile_count == 1 and not odd_k_tile -+ else: -+ is_order_chain_last_tile = k_tile_count == 1 and odd_k_tile - ( - mma_s_consumer_state, - p_mma_producer_state, -@@ -2618,8 +2858,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - row_sum, - correction_factor, - is_second_compute_warp, -+ k_index >= first_mask_tile_idx, - True, -- True, -+ is_order_chain_last_tile, - ) - k_index = k_index + 2 - k_tile_count = k_tile_count - 1 -@@ -2635,14 +2876,17 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - - if odd_k_tile: - if cutlass.const_expr(is_second_compute_warp): -- # next first compute warp in this wave -- p_mma_producer_state.advance() -- mma_s_consumer_state.advance() -- p_cor_producer_state.advance() -- # first compute warp in next wave -+ # Keep g1 exactly one pipeline entry ahead of g0 at the next -+ # logical work item. When this split has one tile, g1 owns no -+ # tile and needs only this single advance. For larger odd -+ # splits it also skips the final g0-owned stage below. - p_mma_producer_state.advance() - mma_s_consumer_state.advance() - p_cor_producer_state.advance() -+ if k_tile_count_init > 1: -+ p_mma_producer_state.advance() -+ mma_s_consumer_state.advance() -+ p_cor_producer_state.advance() - else: - p_mma_producer_state.advance() - mma_s_consumer_state.advance() -@@ -2949,6 +3193,7 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - is_second_compute_warp: bool, - apply_mask: bool, - is_local_last_tile: cutlass.Boolean, -+ is_order_chain_last_tile: cutlass.Boolean, - ) -> tuple[ - pipeline.PipelineState, - pipeline.PipelineState, -@@ -2978,11 +3223,14 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - :param correction_factor: The correction factor - :type correction_factor: cutlass.Float32 - :param apply_mask: Whether the tile needs K-bound / causal masking (Python bool -- for the unmasked/masked bulk loops; runtime cutlass.Boolean for the -- split-KV final iter where mask only applies on the global last tile). -+ for the unmasked/masked bulk loops; runtime cutlass.Boolean for a -+ work-split final tile at the exact per-query-tile mask boundary). - :type apply_mask: bool | cutlass.Boolean - :param is_local_last_tile: Whether the last tile is local - :type is_local_last_tile: cutlass.Boolean -+ :param is_order_chain_last_tile: Whether this group owns the final tile -+ in the logical work item's cross-group order chain -+ :type is_order_chain_last_tile: cutlass.Boolean - - :return: The MMA s consumer state, the P MMA producer state, the P correction producer state, the row max, the row sum, and the correction factor - :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32] -@@ -3028,40 +3276,52 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype) - - row_max_new = row_max -- # Spec-decoding (MTP) causal mask: each row represents one (q_token, head) -- # pair; row r's effective K bound is K - (S_q - 1 - q_tok(r)). -- # With fold factor F = self.fold_sq_ratio (fold_sq=True), the M tile is -- # laid out as [F sub_q_tok][num_heads heads] and there are S_q/F outer -- # chunks indexed by blk_coord[1]: -- # q_tok(r) = blk_coord[1] * F + (r_global // num_heads) -- # r_global = row_in_cta + cluster_idx * (M_tile / cluster_m) -- # When fold_sq=False this reduces to q_tok = blk_coord[1]. For S_q=1 -- # this further reduces to k_bound = K (plain K-bound check). -+ # Spec-decoding (MTP) causal mask. For flattened row -+ # r = q_token * H + q_head, -+ # the usual key < K - S_q + 1 + q_token predicate is equivalent to -+ # H * (key - K + S_q) <= r. -+ # With cyclic DCP, local key k maps to W*k+rank, producing -+ # H * (W*k + rank - G + S_q) <= r. -+ # Both forms avoid integer division/modulo. The enclosing apply_mask -+ # branch remains compile-time false for dense bulk K tiles, so they do -+ # not execute any of this row-dependent arithmetic. - # Masked positions are filled with a large negative sentinel (not -inf) -- # to avoid NaN propagation when an entire row becomes masked. -+ # to avoid NaN propagation when an entire row becomes masked. The DCP -+ # epilogue turns such an empty row into the neutral O=0, LSE=-inf pair. - cta_m_rows = self.mma_qk_tiler[0] // self.cluster_shape_mnk[0] - arch = BaseDSL._get_dsl().get_arch_enum() - if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f): - cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc) - for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): - if apply_mask: -- if cutlass.const_expr(self.fold_sq): -- q_tok = ( -- common_params.blk_coord[1] * self.fold_sq_ratio -- + (tTR_tS[i][0] + common_params.blk_coord[0] * cta_m_rows) -- // self.num_heads -+ flat_q_row = ( -+ common_params.blk_coord[1] * self.mma_qk_tiler[0] -+ + common_params.blk_coord[0] * cta_m_rows -+ + tTR_tS[i][0] -+ ) -+ key_pos = tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index -+ if cutlass.const_expr(self.enable_dcp): -+ mask_threshold = self.num_heads * ( -+ self.cp_world * key_pos -+ + common_params.cp_rank -+ - common_params.causal_seq_len -+ + self.seq_len_q -+ ) -+ tTR_rAcc[i] = ( -+ tTR_rAcc[i] -+ if cute.elem_less(key_pos, common_params.K) -+ and not cute.elem_less(flat_q_row, mask_threshold) -+ else self.acc_dtype(-1.0e6) - ) - else: -- q_tok = common_params.blk_coord[1] -- k_bound = common_params.K - (self.seq_len_q - 1) + q_tok -- tTR_rAcc[i] = ( -- tTR_rAcc[i] -- if cute.elem_less( -- tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, -- k_bound, -+ mask_threshold = self.num_heads * ( -+ key_pos - common_params.K + self.seq_len_q -+ ) -+ tTR_rAcc[i] = ( -+ tTR_rAcc[i] -+ if not cute.elem_less(flat_q_row, mask_threshold) -+ else self.acc_dtype(-1.0e6) - ) -- else self.acc_dtype(-1.0e6) -- ) - # reduction for row_max - row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) - elif cutlass.const_expr( -@@ -3092,23 +3352,34 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout) - if apply_mask: - for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): -- if cutlass.const_expr(self.fold_sq): -- q_tok = ( -- common_params.blk_coord[1] * self.fold_sq_ratio -- + (tTR_tS[i][0] + common_params.blk_coord[0] * cta_m_rows) -- // self.num_heads -+ flat_q_row = ( -+ common_params.blk_coord[1] * self.mma_qk_tiler[0] -+ + common_params.blk_coord[0] * cta_m_rows -+ + tTR_tS[i][0] -+ ) -+ key_pos = tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index -+ if cutlass.const_expr(self.enable_dcp): -+ mask_threshold = self.num_heads * ( -+ self.cp_world * key_pos -+ + common_params.cp_rank -+ - common_params.causal_seq_len -+ + self.seq_len_q -+ ) -+ tTR_rAcc[i] = ( -+ tTR_rAcc[i] -+ if cute.elem_less(key_pos, common_params.K) -+ and not cute.elem_less(flat_q_row, mask_threshold) -+ else self.acc_dtype(-1.0e6) - ) - else: -- q_tok = common_params.blk_coord[1] -- k_bound = common_params.K - (self.seq_len_q - 1) + q_tok -- tTR_rAcc[i] = ( -- tTR_rAcc[i] -- if cute.elem_less( -- tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, -- k_bound, -+ mask_threshold = self.num_heads * ( -+ key_pos - common_params.K + self.seq_len_q -+ ) -+ tTR_rAcc[i] = ( -+ tTR_rAcc[i] -+ if not cute.elem_less(flat_q_row, mask_threshold) -+ else self.acc_dtype(-1.0e6) - ) -- else self.acc_dtype(-1.0e6) -- ) - # reduction for row_max after manual masking - row_max_new = tTR_rAcc.load().reduce( - cute.ReductionOp.MAX, row_max_new, 0 -@@ -3152,7 +3423,7 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - # the .arrive() it issues at the bottom of its prev iter, so this - # arrive_and_wait gives us the acquire memory ordering needed for the - # load_other_group_metadata read below. g0's first wait is satisfied -- # by g1's pre-arrive of bar_0 at warp setup (init-phase trick). -+ # by g1's per-work pre-arrive of bar_0 (init-phase trick). - if cutlass.const_expr(is_second_compute_warp): - self.softmax_order_bar_1.arrive_and_wait() - else: -@@ -3285,9 +3556,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - ) - - # Pingpong B — signal the OTHER group to start its critical section. -- # Split-phase .arrive() contributes this group's threads to the other -- # group's bar and falls through immediately. cfence: control-flow fence -- # to prevent ptxas from hoisting subsequent work above bar.arrive. -+ # Intermediate tiles use split-phase arrive so the groups overlap. The -+ # final owner waits for the peer's parity-tail rendezvous, preventing a -+ # persistent group from arriving on the same named-barrier generation -+ # again in the next logical work item before this one has reset. - - # split kv case - if is_local_last_tile: -@@ -3304,9 +3576,15 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - ) - # cute.nvgpu.cfence() - if cutlass.const_expr(is_second_compute_warp): -- self.softmax_order_bar_0.arrive() # g1 → g0 -+ if is_order_chain_last_tile: -+ self.softmax_order_bar_0.arrive_and_wait() # g1 → g0, final even tile -+ else: -+ self.softmax_order_bar_0.arrive() # g1 → g0 - else: -- self.softmax_order_bar_1.arrive() # g0 → g1 -+ if is_order_chain_last_tile: -+ self.softmax_order_bar_1.arrive_and_wait() # g0 → g1, final odd tile -+ else: -+ self.softmax_order_bar_1.arrive() # g0 → g1 - # cute.nvgpu.cfence() - - mma_s_consumer_state.advance() -@@ -3547,6 +3825,24 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - - return mma_o_consumer_state - -+ @cute.jit -+ def dcp_split_has_valid_key( -+ self, -+ common_params: SimpleNamespace, -+ flat_q_row: cutlass.Int32, -+ ) -> cutlass.Boolean: -+ """Whether this row has a visible key in the current local K split.""" -+ first_local_key = common_params.k_index * self.mma_qk_tiler[1] -+ mask_threshold = self.num_heads * ( -+ self.cp_world * first_local_key -+ + common_params.cp_rank -+ - common_params.causal_seq_len -+ + self.seq_len_q -+ ) -+ return cute.elem_less(first_local_key, common_params.K) and not cute.elem_less( -+ flat_q_row, mask_threshold -+ ) -+ - @cute.jit - def epilogue( - self, -@@ -3600,14 +3896,33 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) - - # apply output scale and normalize by row_sum -- for i in cutlass.range( -- cute.size(tTR_rAcc), vectorize=True, unroll_full=True -- ): -- tTR_rAcc[i] = ( -- tTR_rAcc[i] -- * epilogue_params.output_scale -- * cute.arch.rcp_approx(row_sum) -+ row_has_valid_key = True -+ if cutlass.const_expr(self.enable_dcp): -+ flat_q_row = ( -+ common_params.blk_coord[1] * self.mma_qk_tiler[0] + tTR_cO[0][0] -+ ) -+ row_has_valid_key = self.dcp_split_has_valid_key( -+ common_params, flat_q_row -+ ) -+ if cutlass.const_expr(self.enable_dcp): -+ output_normalizer = ( -+ epilogue_params.output_scale * cute.arch.rcp_approx(row_sum) -+ if row_has_valid_key -+ else self.acc_dtype(0.0) - ) -+ for i in cutlass.range( -+ cute.size(tTR_rAcc), vectorize=True, unroll_full=True -+ ): -+ tTR_rAcc[i] = tTR_rAcc[i] * output_normalizer -+ else: -+ for i in cutlass.range( -+ cute.size(tTR_rAcc), vectorize=True, unroll_full=True -+ ): -+ tTR_rAcc[i] = ( -+ tTR_rAcc[i] -+ * epilogue_params.output_scale -+ * cute.arch.rcp_approx(row_sum) -+ ) - - # store o to global memory - tR2G_rO_src = None -@@ -3688,6 +4003,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - cute.math.log2(row_sum, fastmath=True) - + epilogue_params.softmax_scale_log2 * row_max - ) -+ if cutlass.const_expr(self.enable_dcp): -+ lse = lse if row_has_valid_key else -self.lse_dtype.inf - # When writing directly to the user-facing mLSE (single-tile, - # no split-KV merge), convert from log2 base to natural log. - # When writing the per-split intermediate (mAccLSE branch), keep -@@ -3905,25 +4222,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - - return tile_sched_params, grid - -- @staticmethod -- def compute_fold_sq_ratio(num_heads: int, seq_len_q: int, m_tile: int) -> int: -- """Derive the seq_len_q-into-heads fold factor F. -- -- Returns the largest integer F such that: -- - F divides seq_len_q evenly -- - num_heads * F ≤ m_tile -- - 1 ≤ F ≤ seq_len_q -- -- F=1 means no folding (i.e. ``fold_sq`` should be False at the caller). -- """ -- if num_heads >= m_tile: -- return 1 -- max_fold = min(seq_len_q, m_tile // num_heads) -- for f in range(max_fold, 0, -1): -- if seq_len_q % f == 0: -- return f -- return 1 -- - @staticmethod - def get_workspace_size( - H: int, -@@ -3953,12 +4251,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - """ - if split_kv == 1: - return 0 -- # Decode packs heads into a physical 128-wide MMA-M tile. For H < 128, -- # split-KV partials can still touch the padded head lanes before -- # reduction, so size the workspace for max(H, 128). Mirrors the same -- # padding applied in initialize_workspace(). See #3235. -- workspace_heads = max(H, 128) -- return B * workspace_heads * S * split_kv * (D + 1) * acc_dtype.width // 8 -+ # The first workspace mode is the physical MMA-M row extent. Direct -+ # callers may still pass a logical H < 128, so preserve the M128 floor -+ # used by initialize_workspace(). The flat internal path passes 128. -+ workspace_rows = max(H, 128) -+ return B * workspace_rows * S * split_kv * (D + 1) * acc_dtype.width // 8 - - @cute.jit - def initialize_workspace( -@@ -3994,29 +4291,29 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - """ - acc_o, acc_lse = None, None - if cutlass.const_expr(workspace is not None): -- # Pad head dim to the physical 128-wide MMA-M tile. Without this, -- # H<128 split-KV partials write past the workspace. See #3235. -- workspace_H = cutlass.max(H, cutlass.Int32(128)) -+ # The first workspace mode is the physical MMA-M row extent. -+ # Preserve the M128 floor for direct callers that pass H < 128. -+ workspace_rows = cutlass.max(H, cutlass.Int32(128)) - align = 256 // self.q_dtype.width - acc_o_layout = cute.make_layout( -- (workspace_H, split_kv, D, S, B), -+ (workspace_rows, split_kv, D, S, B), - stride=( - cute.assume(split_kv * D, align), - cute.assume(D, align), - 1, -- cute.assume(split_kv * workspace_H * D, align), -- cute.assume(workspace_H * split_kv * S * D, align), -+ cute.assume(split_kv * workspace_rows * D, align), -+ cute.assume(workspace_rows * split_kv * S * D, align), - ), - ) - acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=acc_dtype) - acc_o = cute.make_tensor(acc_o_iter, acc_o_layout) - acc_lse_layout = cute.make_layout( -- (workspace_H, split_kv, S, B), -+ (workspace_rows, split_kv, S, B), - stride=( - split_kv, - 1, -- workspace_H * split_kv, -- workspace_H * split_kv * S, -+ workspace_rows * split_kv, -+ workspace_rows * split_kv * S, - ), - ) - acc_lse_iter = cute.recast_ptr( -@@ -4098,11 +4395,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: - return False - if is_var_split_kv and not is_var_seq: - return False -- if mma_qk_tiler_mn[0] < H: -+ if H <= 0 or mma_qk_tiler_mn[0] < H: - return False -- # When H < M tile, fold up to F tokens of S into H (M_eff = H*F ≤ M_tile). -- # F is auto-picked as the largest divisor of S with H*F ≤ M_tile. -- # F=1 always works, so any (H ≤ M_tile, S ≥ 1) is implementable. -+ # Query-token/head rows are flattened across 128-row M tiles. A -+ # partial final tile is supported, so any H <= M_tile and S >= 1 is -+ # valid. - if S < 1: - return False - if K <= 0: -diff --git a/flashinfer/cute_dsl/attention/monolithic/mla_helpers.py b/flashinfer/cute_dsl/attention/monolithic/mla_helpers.py -index 722bcfdf..0acd6919 100644 ---- a/flashinfer/cute_dsl/attention/monolithic/mla_helpers.py -+++ b/flashinfer/cute_dsl/attention/monolithic/mla_helpers.py -@@ -302,3 +302,34 @@ MAX_SPLITS = 256 - - def ceil_div(a: int, b: int) -> int: - return (a + b - 1) // b -+ -+ -+def compute_q_tile_layout( -+ num_heads: int, seq_len_q: int, m_tile: int = 128 -+) -> tuple[int, int, int]: -+ """Return ``(total_rows, num_tiles, tail_rows)`` for flat query packing. -+ -+ Query-token and head modes are one affine row space ordered as -+ ``flat_row = q_token * num_heads + q_head``. Consecutive M tiles may -+ therefore cross token boundaries; only the final tile can be partial. -+ ``tail_rows`` is always in ``[1, m_tile]`` and equals ``m_tile`` when the -+ flattened row count exactly fills the final tile. -+ -+ This host-side helper is shared by launch/workspace selection and both -+ kernel variants so their split-KV geometry cannot drift apart. -+ """ -+ if num_heads <= 0: -+ raise ValueError(f"num_heads must be positive, got {num_heads}") -+ if seq_len_q <= 0: -+ raise ValueError(f"seq_len_q must be positive, got {seq_len_q}") -+ if m_tile <= 0: -+ raise ValueError(f"m_tile must be positive, got {m_tile}") -+ if num_heads > m_tile: -+ raise ValueError( -+ f"num_heads ({num_heads}) must not exceed the MMA M tile ({m_tile})" -+ ) -+ -+ total_rows = num_heads * seq_len_q -+ num_tiles = ceil_div(total_rows, m_tile) -+ tail_rows = total_rows - (num_tiles - 1) * m_tile -+ return total_rows, num_tiles, tail_rows -diff --git a/flashinfer/mla/_core.py b/flashinfer/mla/_core.py -index bf422ec6..84b2c11f 100644 ---- a/flashinfer/mla/_core.py -+++ b/flashinfer/mla/_core.py -@@ -1983,13 +1983,66 @@ def _round_to_seq_len_bucket(x: int) -> int: - return 1 << (x - 1).bit_length() - - -+def _resolve_cute_dsl_workspace_sizer( -+ cute_dsl_impl: str, -+ sinks: Optional[Union[List[torch.Tensor], Tuple[torch.Tensor, ...], torch.Tensor]], -+ enable_dcp: bool = False, -+): -+ """Resolve the selected CuTeDSL implementation and its workspace policy.""" -+ from ..cute_dsl.attention.mla_dispatch import _resolve_impl -+ -+ resolved_impl = _resolve_impl( -+ requested=cute_dsl_impl, -+ kwargs={"sinks": sinks, "enable_dcp": enable_dcp}, -+ ) -+ if resolved_impl == "monolithic": -+ from ..cute_dsl.attention.monolithic.mla_decode import ( -+ _get_split_kv_and_workspace_size, -+ ) -+ else: -+ from ..cute_dsl.attention.wrappers.batch_mla import ( -+ _get_split_kv_and_workspace_size, -+ ) -+ return _get_split_kv_and_workspace_size, resolved_impl -+ -+ -+def _get_cute_dsl_workspace_sizer( -+ cute_dsl_impl: str, -+ sinks: Optional[Union[List[torch.Tensor], Tuple[torch.Tensor, ...], torch.Tensor]], -+ enable_dcp: bool = False, -+): -+ """Return the workspace policy owned by the selected CuTeDSL implementation.""" -+ return _resolve_cute_dsl_workspace_sizer(cute_dsl_impl, sinks, enable_dcp)[0] -+ -+ -+def _call_cute_dsl_workspace_sizer( -+ workspace_sizer, -+ resolved_impl: str, -+ batch_size: int, -+ q_len: int, -+ num_heads: int, -+ kv_lora_rank: int, -+ max_active_blocks: int, -+ max_seq_len: int, -+): -+ """Call an implementation's workspace policy with its supported arguments.""" -+ args = (batch_size, q_len, num_heads, kv_lora_rank, max_active_blocks) -+ if resolved_impl == "monolithic": -+ return workspace_sizer(*args, max_seq_len) -+ return workspace_sizer(*args) -+ -+ - def _cute_dsl_max_supported_batch( - workspace_bytes: int, - q_len: int, - num_heads: int, - kv_lora_rank: int, - max_active_blocks: int, -+ max_seq_len: int, - candidate_max: int, -+ cute_dsl_impl: str, -+ sinks: Optional[Union[List[torch.Tensor], Tuple[torch.Tensor, ...], torch.Tensor]], -+ enable_dcp: bool = False, - ) -> int: - """Largest batch the caller's workspace can support for cute-dsl MLA decode. - -@@ -1997,15 +2050,22 @@ def _cute_dsl_max_supported_batch( - split-K state. Binary-search for the largest ``B <= candidate_max`` whose - ``get_workspace_size(...)`` fits in ``workspace_bytes``. - """ -- from ..cute_dsl.attention.wrappers.batch_mla import ( -- _get_split_kv_and_workspace_size, -+ workspace_sizer, resolved_impl = _resolve_cute_dsl_workspace_sizer( -+ cute_dsl_impl, sinks, enable_dcp - ) - - lo, hi = 1, max(1, candidate_max) - while lo < hi: - mid = (lo + hi + 1) // 2 -- _, ws = _get_split_kv_and_workspace_size( -- mid, q_len, num_heads, kv_lora_rank, max_active_blocks -+ _, ws = _call_cute_dsl_workspace_sizer( -+ workspace_sizer, -+ resolved_impl, -+ mid, -+ q_len, -+ num_heads, -+ kv_lora_rank, -+ max_active_blocks, -+ max_seq_len, - ) - if ws <= workspace_bytes: - lo = mid -@@ -2020,7 +2080,11 @@ def _compute_mla_decode_buckets( - q_len: int, - num_heads: int, - kv_lora_rank: int, -+ max_seq_len: int, - device: torch.device, -+ cute_dsl_impl: str, -+ sinks: Optional[Union[List[torch.Tensor], Tuple[torch.Tensor, ...], torch.Tensor]], -+ enable_dcp: bool = False, - ) -> Tuple[int, ...]: - """Compute the autotune bucket list from kernel/workspace limits only. - -@@ -2048,13 +2112,112 @@ def _compute_mla_decode_buckets( - num_heads=num_heads, - kv_lora_rank=kv_lora_rank, - max_active_blocks=get_num_sm(device), -+ max_seq_len=max_seq_len, - candidate_max=_TRTLLM_GEN_MLA_MAX_BATCH, -+ cute_dsl_impl=cute_dsl_impl, -+ sinks=sinks, -+ enable_dcp=enable_dcp, - ) - cap = max(cap, cute_dsl_cap) - - return get_hybrid_num_tokens_buckets(max(1, cap)) - - -+def _validate_mla_dcp_args( -+ *, -+ query: torch.Tensor, -+ backend: str, -+ sinks: Optional[List[torch.Tensor]], -+ cum_seq_lens_q: Optional[torch.Tensor], -+ max_q_len: Optional[int], -+ return_lse: bool, -+ enable_dcp: bool, -+ cp_world: int, -+ cp_rank: int, -+ causal_seqlens_kv_global: Optional[torch.Tensor], -+) -> str: -+ """Validate the public DCP contract and return the effective backend.""" -+ if not isinstance(enable_dcp, bool): -+ raise TypeError(f"enable_dcp must be a bool, got {type(enable_dcp).__name__}") -+ if not isinstance(cp_world, int) or isinstance(cp_world, bool) or cp_world <= 0: -+ raise ValueError(f"cp_world must be a positive integer, got {cp_world!r}") -+ if not isinstance(cp_rank, int) or isinstance(cp_rank, bool): -+ raise TypeError(f"cp_rank must be an integer, got {type(cp_rank).__name__}") -+ -+ if not enable_dcp: -+ nondefault = [] -+ if cp_world != 1: -+ nondefault.append(f"cp_world={cp_world}") -+ if cp_rank != 0: -+ nondefault.append(f"cp_rank={cp_rank}") -+ if causal_seqlens_kv_global is not None: -+ nondefault.append("causal_seqlens_kv_global") -+ if nondefault: -+ raise ValueError( -+ "DCP arguments require enable_dcp=True; got " + ", ".join(nondefault) -+ ) -+ return backend -+ -+ if query.ndim != 4: -+ raise ValueError( -+ "DCP requires a dense query with shape " -+ "[batch_size, q_len_per_request, num_heads, head_dim_qk]" -+ ) -+ if not 0 <= cp_rank < cp_world: -+ raise ValueError( -+ f"cp_rank must satisfy 0 <= cp_rank < cp_world, got " -+ f"cp_rank={cp_rank}, cp_world={cp_world}" -+ ) -+ if backend not in ("auto", "cute-dsl"): -+ raise ValueError( -+ f"enable_dcp=True is only supported by backend='cute-dsl', got " -+ f"backend={backend!r}" -+ ) -+ if not return_lse: -+ raise ValueError( -+ "enable_dcp=True requires return_lse=True so rank-local " -+ "attention states can be merged" -+ ) -+ if sinks is not None: -+ raise ValueError( -+ "DCP cannot be combined with sinks: DCP requires monolithic " -+ "CuTeDSL MLA, while sinks require the modular implementation" -+ ) -+ if cum_seq_lens_q is not None or max_q_len is not None: -+ raise ValueError("DCP does not support cum_seq_lens_q / max_q_len") -+ if causal_seqlens_kv_global is None: -+ raise ValueError( -+ "causal_seqlens_kv_global is required when enable_dcp=True" -+ ) -+ if not isinstance(causal_seqlens_kv_global, torch.Tensor): -+ raise TypeError( -+ "causal_seqlens_kv_global must be a torch.Tensor, got " -+ f"{type(causal_seqlens_kv_global).__name__}" -+ ) -+ if causal_seqlens_kv_global.dtype != torch.int32: -+ raise ValueError( -+ "causal_seqlens_kv_global must have dtype torch.int32, got " -+ f"{causal_seqlens_kv_global.dtype}" -+ ) -+ if not causal_seqlens_kv_global.is_cuda: -+ raise ValueError("causal_seqlens_kv_global must be a CUDA tensor") -+ if causal_seqlens_kv_global.device != query.device: -+ raise ValueError( -+ "causal_seqlens_kv_global must be on the query device " -+ f"{query.device}, got {causal_seqlens_kv_global.device}" -+ ) -+ if tuple(causal_seqlens_kv_global.shape) != (query.shape[0],): -+ raise ValueError( -+ "causal_seqlens_kv_global must have shape " -+ f"({query.shape[0]},), got {tuple(causal_seqlens_kv_global.shape)}" -+ ) -+ if not causal_seqlens_kv_global.is_contiguous(): -+ raise ValueError("causal_seqlens_kv_global must be contiguous") -+ -+ # DCP masking exists only in the monolithic CuTeDSL implementation. -+ return "cute-dsl" -+ -+ - def _cute_dsl_incompatibility_reason( - query: torch.Tensor, - out_dtype: torch.dtype, -@@ -2068,9 +2231,9 @@ def _cute_dsl_incompatibility_reason( - kv_lora_rank: int, - page_size: int, - is_var_seq: bool, -- return_lse: bool, -- lse: Optional[torch.Tensor], - cute_dsl_impl: str = "auto", -+ enable_dcp: bool = False, -+ cp_world: int = 1, - ) -> Optional[str]: - """Return None if cute-dsl can handle this call, else a human-readable reason. - -@@ -2123,8 +2286,14 @@ def _cute_dsl_incompatibility_reason( - try: - from ..cute_dsl.attention.mla_dispatch import _resolve_impl - -- resolved_impl = _resolve_impl(requested=cute_dsl_impl, kwargs={"sinks": sinks}) -- except (ValueError, ImportError) as e: -+ resolved_impl = _resolve_impl( -+ requested=cute_dsl_impl, -+ kwargs={ -+ "sinks": sinks, -+ "enable_dcp": enable_dcp, -+ }, -+ ) -+ except (TypeError, ValueError, ImportError) as e: - return f"cute-dsl backend (MLA decode kernel): {e}" - - try: -@@ -2133,7 +2302,7 @@ def _cute_dsl_incompatibility_reason( - else: - from ..cute_dsl.attention.wrappers.batch_mla import _check_can_implement - -- _check_can_implement( -+ check_kwargs = dict( - torch_dtype=query.dtype, - torch_out_dtype=out_dtype, - page_size=page_size, -@@ -2145,7 +2314,10 @@ def _cute_dsl_incompatibility_reason( - is_var_seq=is_var_seq, - is_var_split_kv=False, - ) -- except (ValueError, ImportError) as e: -+ if resolved_impl == "monolithic": -+ check_kwargs.update(enable_dcp=enable_dcp, cp_world=cp_world) -+ _check_can_implement(**check_kwargs) -+ except (TypeError, ValueError, ImportError) as e: - return f"cute-dsl backend (MLA decode kernel) cannot implement this configuration: {e}" - return None - -@@ -2155,8 +2327,11 @@ def _mla_decode_tuning_config( - buckets: tuple[int, ...], - num_pages: int, - profile_seq_len: int, -+ enable_dcp: bool = False, -+ cp_world: int = 1, -+ cp_rank: int = 0, - ) -> TuningConfig: -- """One TuningConfig (and one pair of initializer closures) per key. -+ """One TuningConfig and stable initializer set per key. - - Memoized because ``AutoTuner._find_nearest_profile`` lru-caches on - ``(shapes, tuning_config)``: a fresh config per dispatcher call shares -@@ -2164,12 +2339,13 @@ def _mla_decode_tuning_config( - ``tensor_initializers``) but never compares equal (closures compare by - identity), so would result in a leak. - -- The DynamicTensorSpec sweeps batch dim across all four ``inputs`` tensors -- (query, block_tables, seq_lens, out). ``block_tables`` is initialized via -- ``random_(0, num_pages)`` which wraps mod kv_cache size — safe for autotune -- profiling because MLA decode reads kv_cache and never writes it, so aliased -- page reads give correct timing measurements. ``seq_lens`` is filled -- homogeneously with ``profile_seq_len``. -+ The DynamicTensorSpec sweeps batch dim across ``query``, ``block_tables``, -+ ``seq_lens``, ``out``, and, for DCP, ``causal_seqlens_kv_global``. -+ ``block_tables`` is initialized via ``random_(0, num_pages)`` which wraps -+ mod kv_cache size — safe for autotune profiling because MLA decode reads -+ kv_cache and never writes it, so aliased page reads give correct timing -+ measurements. ``seq_lens`` is filled with ``profile_seq_len``; the -+ synthetic DCP global bound preserves that exact rank-local length. - """ - - def init_block_tables(shapes, dtype, device): -@@ -2182,14 +2358,32 @@ def _mla_decode_tuning_config( - tensor.fill_(profile_seq_len) - return tensor - -+ def init_causal_seqlens_kv_global(shapes, dtype, device): -+ tensor = torch.empty(shapes, dtype=dtype, device=device) -+ tensor.fill_(profile_seq_len * cp_world + cp_rank) -+ return tensor -+ -+ input_idx = (0, 1, 2, 3, 4) if enable_dcp else (0, 1, 2, 3) -+ tensor_initializers = ( -+ ( -+ None, -+ init_block_tables, -+ init_seq_lens, -+ None, -+ init_causal_seqlens_kv_global, -+ ) -+ if enable_dcp -+ else (None, init_block_tables, init_seq_lens, None) -+ ) -+ - return TuningConfig( - dynamic_tensor_specs=( - DynamicTensorSpec( -- input_idx=(0, 1, 2, 3), -- dim_idx=(0, 0, 0, 0), -+ input_idx=input_idx, -+ dim_idx=(0,) * len(input_idx), - gen_tuning_buckets=buckets, - map_to_tuning_buckets=make_bucket_mapper(buckets, round_map=False), -- tensor_initializers=(None, init_block_tables, init_seq_lens, None), -+ tensor_initializers=tensor_initializers, - ), - ), - use_cuda_graph=True, -@@ -2207,6 +2401,13 @@ def _build_mla_decode_tuning_config( - kv_lora_rank: int, - max_seq_len: int, - device: torch.device, -+ cute_dsl_impl: str = "auto", -+ sinks: Optional[ -+ Union[List[torch.Tensor], Tuple[torch.Tensor, ...], torch.Tensor] -+ ] = None, -+ enable_dcp: bool = False, -+ cp_world: int = 1, -+ cp_rank: int = 0, - ) -> TuningConfig: - """Reduce call args to the memoization key of ``_mla_decode_tuning_config``. - -@@ -2229,10 +2430,21 @@ def _build_mla_decode_tuning_config( - q_len, - num_heads, - kv_lora_rank, -+ max_seq_len, - device, -+ cute_dsl_impl, -+ sinks, -+ enable_dcp, - ) - -- return _mla_decode_tuning_config(buckets, num_pages, profile_seq_len) -+ return _mla_decode_tuning_config( -+ buckets, -+ num_pages, -+ profile_seq_len, -+ enable_dcp, -+ cp_world, -+ cp_rank, -+ ) - - - class TrtllmGenMlaDecodeRunner(TunableRunner): -@@ -2458,6 +2670,9 @@ class CuteDslMlaDecodeRunner(TunableRunner): - return_lse: bool, - sinks: Optional[torch.Tensor], - cute_dsl_impl: str, -+ enable_dcp: bool = False, -+ cp_world: int = 1, -+ cp_rank: int = 0, - ): - from ..cute_dsl.attention import cute_dsl_mla_decode - -@@ -2482,6 +2697,13 @@ class CuteDslMlaDecodeRunner(TunableRunner): - self.return_lse = return_lse - self.sinks = sinks - self.cute_dsl_impl = cute_dsl_impl -+ self.enable_dcp = enable_dcp -+ self.cp_world = cp_world -+ self.cp_rank = cp_rank -+ self._profile_lse: Optional[torch.Tensor] = None -+ self._workspace_sizer, self._resolved_cute_dsl_impl = ( -+ _resolve_cute_dsl_workspace_sizer(cute_dsl_impl, sinks, enable_dcp) -+ ) - - def __hash__(self): - # See TrtllmGenMlaDecodeRunner.__hash__ — tactic-determining state is -@@ -2495,15 +2717,19 @@ class CuteDslMlaDecodeRunner(TunableRunner): - # If the caller's workspace can't fit batch=B for this profile, opt - # out so the autotuner skips us (no JIT cost) and trtllm-gen wins by - # default for that bucket. -- from ..cute_dsl.attention.wrappers.batch_mla import ( -- _get_split_kv_and_workspace_size, -- ) - from ..cute_dsl.utils import get_num_sm - - q = inputs[0] - B, q_len, num_heads, _ = q.shape -- _, ws = _get_split_kv_and_workspace_size( -- B, q_len, num_heads, self.kv_lora_rank, get_num_sm(q.device) -+ _, ws = _call_cute_dsl_workspace_sizer( -+ self._workspace_sizer, -+ self._resolved_cute_dsl_impl, -+ B, -+ q_len, -+ num_heads, -+ self.kv_lora_rank, -+ get_num_sm(q.device), -+ self.max_seq_len, - ) - workspace_bytes = ( - self.workspace_buffer.numel() * self.workspace_buffer.element_size() -@@ -2513,14 +2739,24 @@ class CuteDslMlaDecodeRunner(TunableRunner): - return [-1] - - def get_cache_key_extras(self, inputs): -- q, _, _, out = inputs -+ q, _, _, out = inputs[:4] - # Cute-dsl rejects sparse/skip-softmax/tensor-scales upstream, so - # those are omitted from extras as constants for this runner. - # ``sinks`` and ``cute_dsl_impl`` are included because they flip the -- # impl (modular vs monolithic) inside ``cute_dsl_mla_decode``. -+ # impl (modular vs monolithic) inside ``cute_dsl_mla_decode``. The -+ # monolithic K-tile count and workspace capacity keep cache hits from -+ # bypassing a different split-workspace validity decision. - sinks_key = ( - None if self.sinks is None else (tuple(self.sinks.shape), self.sinks.dtype) - ) -+ workspace_bytes = ( -+ self.workspace_buffer.numel() * self.workspace_buffer.element_size() -+ ) -+ seq_len_workspace_key = ( -+ (self.max_seq_len + 127) // 128 -+ if self._resolved_cute_dsl_impl == "monolithic" -+ else _round_to_seq_len_bucket(self.max_seq_len) -+ ) - return ( - q.dtype, - self.kv_cache.dtype, -@@ -2529,12 +2765,15 @@ class CuteDslMlaDecodeRunner(TunableRunner): - self.kv_lora_rank, - self.qk_rope_head_dim, - self.page_size, -- _round_to_seq_len_bucket(self.max_seq_len), -+ seq_len_workspace_key, -+ workspace_bytes, - self.is_var_seq, - self.uses_shared_paged_kv_idx, - self.enable_pdl, - sinks_key, - self.cute_dsl_impl, -+ getattr(self, "enable_dcp", False), -+ getattr(self, "cp_world", 1), - ) - - def forward( -@@ -2544,7 +2783,31 @@ class CuteDslMlaDecodeRunner(TunableRunner): - do_preparation: bool = False, - **kwargs, - ): -- query, block_tables, seq_lens, out = inputs -+ query, block_tables, seq_lens, out = inputs[:4] -+ causal_seqlens_kv_global = inputs[4] if self.enable_dcp else None -+ -+ # LSE is not a tuning input because it does not influence tactic -+ # selection. When a synthetic batch differs from the caller batch, -+ # provide matching temporary storage while retaining the caller's LSE -+ # for the final invocation. -+ lse = self.lse -+ if self.return_lse: -+ expected_numel = query.shape[0] * query.shape[1] * query.shape[2] -+ if lse is None or lse.numel() != expected_numel: -+ expected_shape = ( -+ query.shape[0] * query.shape[1], -+ query.shape[2], -+ ) -+ if ( -+ self._profile_lse is None -+ or tuple(self._profile_lse.shape) != expected_shape -+ ): -+ self._profile_lse = torch.empty( -+ expected_shape, -+ dtype=torch.float32, -+ device=query.device, -+ ) -+ lse = self._profile_lse - return self._run( - query=query, - kv_cache=self.kv_cache, -@@ -2560,10 +2823,14 @@ class CuteDslMlaDecodeRunner(TunableRunner): - out_dtype=self.out_dtype, - is_var_seq=self.is_var_seq, - enable_pdl=self.enable_pdl, -- lse=self.lse, -+ lse=lse, - return_lse=self.return_lse, - sinks=self.sinks, - cute_dsl_impl=self.cute_dsl_impl, -+ enable_dcp=self.enable_dcp, -+ cp_world=self.cp_world, -+ cp_rank=self.cp_rank, -+ causal_seqlens_kv_global=causal_seqlens_kv_global, - ) - - -@@ -2595,6 +2862,10 @@ def trtllm_batch_decode_with_kv_cache_mla( - cum_seq_lens_q: Optional[torch.Tensor] = None, - max_q_len: Optional[int] = None, - multi_ctas_kv_counter_buffer: Optional[torch.Tensor] = None, -+ enable_dcp: bool = False, -+ cp_world: int = 1, -+ cp_rank: int = 0, -+ causal_seqlens_kv_global: Optional[torch.Tensor] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - r"""Decode MLA with TRTLLM-GEN, CuteDSL, XQA, or SM120/SM121 sparse kernels. - -@@ -2639,12 +2910,15 @@ def trtllm_batch_decode_with_kv_cache_mla( - With ``backend="trtllm-gen"``, the final dimension may use its native - width and does not need padding to a multiple of ``128 / page_size``. - seq_lens : Optional[torch.Tensor] -- Per-request KV sequence lengths for dense and TRTLLM-GEN paths. For -+ Per-request physical KV sequence lengths for dense and TRTLLM-GEN -+ paths. With DCP these are rank-local lengths and continue to control -+ paging, memory bounds, and split-KV. For - SM120/SM121 sparse v32/GLM, pass ``[batch_size, q_len_per_request]`` or - flattened ``[batch_size * q_len_per_request]`` active top-k lengths; if - ``None``, every column in ``block_tables`` is active. - max_seq_len : int -- Maximum KV sequence length used for dense/TRTLLM-GEN scheduling. -+ Maximum physical KV sequence length used for dense/TRTLLM-GEN -+ scheduling. With DCP this is the maximum rank-local length. - Ignored by the SM120/SM121 sparse v32/GLM backend. - sparse_mla_top_k : int - Enables sparse MLA when greater than zero. On SM100/SM103 this selects -@@ -2756,6 +3030,18 @@ def trtllm_batch_decode_with_kv_cache_mla( - for each concurrently executing CUDA stream or graph. Autotune profiling - uses runner-owned internal storage; the caller buffer is used only for the - final request. -+ enable_dcp : bool = False -+ Statically enable cyclic decode context parallelism in the monolithic -+ CuTeDSL MLA kernel. DCP returns a rank-local output/LSE state, so -+ ``return_lse=True`` is required and the caller must merge rank states. -+ cp_world : int = 1 -+ Compile-time context-parallel world size. Rank ``r`` stores global KV -+ positions whose token index modulo ``cp_world`` equals ``r``. -+ cp_rank : int = 0 -+ Runtime-uniform context-parallel rank. -+ causal_seqlens_kv_global : Optional[torch.Tensor] = None -+ Contiguous CUDA int32 tensor ``[batch_size]`` containing the global -+ exclusive causal bound for the newest query token. Required with DCP. - - Note - ---- -@@ -2779,11 +3065,12 @@ def trtllm_batch_decode_with_kv_cache_mla( - Autotune - -------- - On SM100/SM103 dense MLA, calling under ``flashinfer.autotune(True)`` with -- ``backend="auto"`` profiles both ``trtllm-gen`` and ``cute-dsl`` across a -- bucketed batch sweep up to each runner's kernel/workspace cap and caches the -- winning runner per shape signature. Subsequent calls under -- ``autotune(False)`` dispatch to the cached choice; any batch outside the -- tuned range falls back to a default runner with a one-time warning. -+ ``backend="auto"`` profiles both ``trtllm-gen`` and ``cute-dsl`` when DCP -+ is disabled. DCP profiles only its required monolithic CuTeDSL runner. -+ Both modes use a bucketed batch sweep up to each runner's kernel/workspace -+ cap and cache the winning tactic per shape signature. Subsequent calls under -+ ``autotune(False)`` use the cached choice; any batch outside the tuned range -+ falls back to the runner's default tactic with a one-time warning. - - The autotune bucket range and cache key do **not** depend on - ``kv_cache.shape[0]`` (the number of pages in the pool), so reallocating the -@@ -2805,6 +3092,19 @@ def trtllm_batch_decode_with_kv_cache_mla( - if max_q_len is not None and cum_seq_lens_q is None: - raise ValueError("max_q_len is only supported when cum_seq_lens_q is provided") - -+ backend = _validate_mla_dcp_args( -+ query=query, -+ backend=backend, -+ sinks=sinks, -+ cum_seq_lens_q=cum_seq_lens_q, -+ max_q_len=max_q_len, -+ return_lse=return_lse, -+ enable_dcp=enable_dcp, -+ cp_world=cp_world, -+ cp_rank=cp_rank, -+ causal_seqlens_kv_global=causal_seqlens_kv_global, -+ ) -+ - if backend == "auto": - cc = get_compute_capability(query.device) - if cc[0] == 12 and sparse_mla_top_k > 0: -@@ -3116,9 +3416,9 @@ def trtllm_batch_decode_with_kv_cache_mla( - kv_lora_rank, - page_size, - is_var_seq, -- return_lse, -- lse, - cute_dsl_impl, -+ enable_dcp, -+ cp_world, - ) - if backend == "cute-dsl": - if cute_dsl_reason is not None: -@@ -3203,6 +3503,9 @@ def trtllm_batch_decode_with_kv_cache_mla( - return_lse=return_lse, - sinks=cute_dsl_sinks, - cute_dsl_impl=cute_dsl_impl, -+ enable_dcp=enable_dcp, -+ cp_world=cp_world, -+ cp_rank=cp_rank, - ) - ) - -@@ -3217,8 +3520,17 @@ def trtllm_batch_decode_with_kv_cache_mla( - kv_lora_rank=kv_lora_rank, - max_seq_len=max_seq_len, - device=query.device, -+ cute_dsl_impl=cute_dsl_impl, -+ sinks=sinks, -+ enable_dcp=enable_dcp, -+ cp_world=cp_world, -+ cp_rank=cp_rank, - ) - inputs = [query, block_tables, seq_lens, out] -+ if enable_dcp: -+ # The global causal bound varies with batch and must be synthesized -+ # alongside the other batch-shaped tensors during autotuning. -+ inputs.append(causal_seqlens_kv_global) - runner, tactic = AutoTuner.get().choose_one( - "trtllm_batch_decode_mla", - runners, -diff --git a/flashinfer/trace/templates/attention.py b/flashinfer/trace/templates/attention.py -index 610afb82..7b32bb88 100644 ---- a/flashinfer/trace/templates/attention.py -+++ b/flashinfer/trace/templates/attention.py -@@ -2692,6 +2692,11 @@ trtllm_batch_decode_mla_sparse_trace = TraceTemplate( - - - def trtllm_batch_decode_mla_trace_dispatch(**kwargs): -+ if kwargs.get("enable_dcp", False): -+ raise NotImplementedError( -+ "fi_trace does not yet represent cyclic DCP KV ownership or " -+ "cross-rank LSE merging for MLA decode" -+ ) - sparse_mla_top_k = int(kwargs.get("sparse_mla_top_k", 0) or 0) - if sparse_mla_top_k > 0: - return trtllm_batch_decode_mla_sparse_trace -diff --git a/tests/attention/test_cute_dsl_mla_dcp.py b/tests/attention/test_cute_dsl_mla_dcp.py -new file mode 100644 -index 00000000..ee1bfabe ---- /dev/null -+++ b/tests/attention/test_cute_dsl_mla_dcp.py -@@ -0,0 +1,1126 @@ -+# Copyright (c) 2026 by FlashInfer team. -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+# You may obtain a copy of the License at -+# -+# http://www.apache.org/licenses/LICENSE-2.0 -+# -+# Unless required by applicable law or agreed to in writing, software -+# distributed under the License is distributed on an "AS IS" BASIS, -+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -+# See the License for the specific language governing permissions and -+# limitations under the License. -+ -+"""Correctness tests for static DCP support in monolithic CuTe DSL MLA.""" -+ -+import math -+from bisect import bisect_left -+from unittest import mock -+ -+import pytest -+import torch -+ -+from flashinfer.cute_dsl import is_cute_dsl_available -+from flashinfer.utils import is_sm100a_supported, is_sm110a_supported -+ -+ -+_LATENT_DIM = 512 -+_ROPE_DIM = 64 -+_QK_DIM = _LATENT_DIM + _ROPE_DIM -+_PAGE_SIZE = 64 -+ -+ -+def _skip_if_unsupported() -> None: -+ device = torch.device("cuda") -+ if not (is_sm100a_supported(device) or is_sm110a_supported(device)): -+ pytest.skip("Requires SM100-SM110 (tcgen05)") -+ if not is_cute_dsl_available(): -+ pytest.skip("CuTe DSL not available") -+ -+ -+def _ceil_div(numerator: int, denominator: int) -> int: -+ return -(-numerator // denominator) -+ -+ -+def _local_causal_bound( -+ global_bound_newest: int, -+ q_len: int, -+ q_idx: int, -+ cp_world: int, -+ cp_rank: int, -+) -> int: -+ return _ceil_div( -+ global_bound_newest - cp_rank - (q_len - 1) + q_idx, -+ cp_world, -+ ) -+ -+ -+def _flat_dcp_score_is_valid( -+ flat_q_row: int, -+ num_heads: int, -+ local_key: int, -+ q_len: int, -+ global_bound_newest: int, -+ cp_world: int, -+ cp_rank: int, -+) -> bool: -+ return flat_q_row >= num_heads * ( -+ cp_world * local_key -+ + cp_rank -+ - global_bound_newest -+ + q_len -+ ) -+ -+ -+def _local_length(global_length: int, cp_world: int, cp_rank: int) -> int: -+ return max(_ceil_div(global_length - cp_rank, cp_world), 0) -+ -+ -+def _make_inputs( -+ *, -+ global_length: int, -+ q_len: int, -+ num_heads: int, -+ dtype: torch.dtype, -+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: -+ """Create one batch of quantized-once Q and global-coordinate KV.""" -+ return _make_batched_inputs( -+ global_lengths=(global_length,), -+ q_len=q_len, -+ num_heads=num_heads, -+ dtype=dtype, -+ ) -+ -+ -+def _make_batched_inputs( -+ *, -+ global_lengths: tuple[int, ...], -+ q_len: int, -+ num_heads: int, -+ dtype: torch.dtype, -+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: -+ """Create a heterogeneous batch with one padded global-coordinate KV pool.""" -+ device = torch.device("cuda") -+ storage_dtype = torch.float16 if dtype == torch.float8_e4m3fn else dtype -+ batch_size = len(global_lengths) -+ query = ( -+ torch.randn( -+ batch_size, -+ q_len, -+ num_heads, -+ _QK_DIM, -+ device=device, -+ dtype=storage_dtype, -+ ) -+ * 0.1 -+ ).to(dtype) -+ global_kv = ( -+ torch.randn( -+ batch_size, -+ max(global_lengths), -+ _QK_DIM, -+ device=device, -+ dtype=storage_dtype, -+ ) -+ * 0.1 -+ ).to(dtype) -+ global_lens = torch.tensor( -+ global_lengths, dtype=torch.int32, device=device -+ ) -+ return query, global_kv, global_lens -+ -+ -+def _pack_cyclic_rank_pages( -+ global_kv: torch.Tensor, -+ global_lens: torch.Tensor, -+ cp_world: int, -+ cp_rank: int, -+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: -+ """Pack ``g % cp_world == cp_rank`` tokens into a contiguous paged cache.""" -+ batch_size, _, d_qk = global_kv.shape -+ global_lens_host = global_lens.tolist() -+ local_lens_host = [ -+ _local_length(global_len, cp_world, cp_rank) -+ for global_len in global_lens_host -+ ] -+ pages_per_batch = [ -+ max(1, _ceil_div(local_len, _PAGE_SIZE)) -+ for local_len in local_lens_host -+ ] -+ max_pages = max(pages_per_batch) -+ total_pages = sum(pages_per_batch) -+ -+ local_cache = torch.zeros( -+ total_pages, -+ _PAGE_SIZE, -+ d_qk, -+ dtype=global_kv.dtype, -+ device=global_kv.device, -+ ) -+ block_tables = torch.zeros( -+ batch_size, -+ max_pages, -+ dtype=torch.int32, -+ device=global_kv.device, -+ ) -+ -+ next_page = 0 -+ for batch_idx, (global_len, local_len, num_pages) in enumerate( -+ zip( -+ global_lens_host, -+ local_lens_host, -+ pages_per_batch, -+ strict=True, -+ ) -+ ): -+ page_ids = torch.arange( -+ next_page, -+ next_page + num_pages, -+ dtype=torch.int32, -+ device=global_kv.device, -+ ) -+ block_tables[batch_idx, :num_pages] = page_ids -+ if local_len: -+ local_tokens = global_kv[ -+ batch_idx, cp_rank:global_len:cp_world -+ ] -+ local_cache[next_page : next_page + num_pages].view( -+ -1, d_qk -+ )[:local_len].copy_(local_tokens) -+ next_page += num_pages -+ -+ local_lens = torch.tensor( -+ local_lens_host, dtype=torch.int32, device=global_kv.device -+ ) -+ return ( -+ local_cache, -+ block_tables, -+ local_lens, -+ max(1, max(local_lens_host)), -+ ) -+ -+ -+def _reference_attention( -+ query: torch.Tensor, -+ global_kv: torch.Tensor, -+ global_lens: torch.Tensor, -+ *, -+ cp_world: int = 1, -+ cp_rank: int = 0, -+) -> tuple[torch.Tensor, torch.Tensor]: -+ """Reference one cyclic rank; world=1 is the full-context reference.""" -+ batch_size, q_len, num_heads, _ = query.shape -+ out = torch.zeros( -+ batch_size, -+ q_len, -+ num_heads, -+ _LATENT_DIM, -+ dtype=torch.float32, -+ device=query.device, -+ ) -+ lse = torch.full( -+ (batch_size, q_len, num_heads), -+ -torch.inf, -+ dtype=torch.float32, -+ device=query.device, -+ ) -+ softmax_scale = 1.0 / math.sqrt(_LATENT_DIM) -+ -+ for batch_idx, global_len in enumerate(global_lens.tolist()): -+ keys = global_kv[ -+ batch_idx, cp_rank:global_len:cp_world -+ ].float() -+ global_positions = range(cp_rank, global_len, cp_world) -+ for q_idx in range(q_len): -+ global_bound = global_len - (q_len - 1) + q_idx -+ # Count directly in global coordinates so this reference remains -+ # independent of the ceiling-divided kernel formula without -+ # synchronizing on a CUDA predicate. -+ visible_count = bisect_left(global_positions, global_bound) -+ if visible_count == 0: -+ continue -+ visible_keys = keys[:visible_count] -+ scores = torch.einsum( -+ "hd,kd->hk", -+ query[batch_idx, q_idx].float(), -+ visible_keys, -+ ) -+ scores *= softmax_scale -+ lse[batch_idx, q_idx] = torch.logsumexp(scores, dim=-1) -+ probabilities = torch.softmax(scores, dim=-1) -+ out[batch_idx, q_idx] = torch.einsum( -+ "hk,kd->hd", -+ probabilities, -+ visible_keys[:, :_LATENT_DIM], -+ ) -+ return out, lse -+ -+ -+def _merge_rank_outputs_natural_log( -+ rank_outputs: list[torch.Tensor], -+ rank_lses: list[torch.Tensor], -+) -> tuple[torch.Tensor, torch.Tensor]: -+ """Stable LSE-weighted merge for the kernel's natural-log public LSE.""" -+ outputs = torch.stack([out.float() for out in rank_outputs], dim=0) -+ lses = torch.stack([lse.float() for lse in rank_lses], dim=0) -+ max_lse = lses.max(dim=0).values -+ has_keys = torch.isfinite(max_lse) -+ safe_max = torch.where(has_keys, max_lse, torch.zeros_like(max_lse)) -+ weights = torch.where( -+ torch.isfinite(lses), -+ torch.exp(lses - safe_max.unsqueeze(0)), -+ torch.zeros_like(lses), -+ ) -+ weight_sum = weights.sum(dim=0) -+ safe_sum = torch.where(has_keys, weight_sum, torch.ones_like(weight_sum)) -+ merged_out = ( -+ outputs * weights.unsqueeze(-1) -+ ).sum(dim=0) / safe_sum.unsqueeze(-1) -+ merged_out = torch.where( -+ has_keys.unsqueeze(-1), merged_out, torch.zeros_like(merged_out) -+ ) -+ merged_lse = torch.where( -+ has_keys, -+ safe_max + torch.log(safe_sum), -+ torch.full_like(max_lse, -torch.inf), -+ ) -+ return merged_out, merged_lse -+ -+ -+def _prepare_rank_call( -+ query: torch.Tensor, -+ global_kv: torch.Tensor, -+ global_lens: torch.Tensor, -+ *, -+ cp_world: int, -+ cp_rank: int, -+ is_var_seq: bool = False, -+ enable_pdl: bool = False, -+) -> tuple[dict, int]: -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ _get_split_kv_and_workspace_size, -+ ) -+ from flashinfer.cute_dsl.utils import get_num_sm -+ -+ kv_cache, block_tables, local_lens, max_local_len = ( -+ _pack_cyclic_rank_pages( -+ global_kv, global_lens, cp_world, cp_rank -+ ) -+ ) -+ batch_size, q_len, num_heads, _ = query.shape -+ split_kv, workspace_size = _get_split_kv_and_workspace_size( -+ batch_size, -+ q_len, -+ num_heads, -+ _LATENT_DIM, -+ get_num_sm(query.device), -+ max_local_len, -+ ) -+ workspace = torch.empty( -+ max(workspace_size, 1), dtype=torch.int8, device=query.device -+ ) -+ return ( -+ { -+ "query": query, -+ "kv_cache": kv_cache, -+ "workspace_buffer": workspace, -+ "kv_lora_rank": _LATENT_DIM, -+ "qk_rope_head_dim": _ROPE_DIM, -+ "block_tables": block_tables, -+ "seq_lens": local_lens, -+ "max_seq_len": max_local_len, -+ "softmax_scale": 1.0 / math.sqrt(_LATENT_DIM), -+ "is_var_seq": is_var_seq, -+ "enable_pdl": enable_pdl, -+ }, -+ split_kv, -+ ) -+ -+ -+def _launch_rank( -+ query: torch.Tensor, -+ global_kv: torch.Tensor, -+ global_lens: torch.Tensor, -+ *, -+ cp_world: int, -+ cp_rank: int, -+ enable_dcp: bool, -+ causal_lens: torch.Tensor | None = None, -+ is_var_seq: bool = False, -+ enable_pdl: bool = False, -+) -> tuple[torch.Tensor, torch.Tensor, int]: -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ cute_dsl_mla_decode, -+ ) -+ -+ call_args, split_kv = _prepare_rank_call( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=cp_world, -+ cp_rank=cp_rank, -+ is_var_seq=is_var_seq, -+ enable_pdl=enable_pdl, -+ ) -+ kwargs = {} -+ if enable_dcp: -+ kwargs = { -+ "enable_dcp": True, -+ "cp_world": cp_world, -+ "cp_rank": cp_rank, -+ "causal_seqlens_kv_global": ( -+ global_lens if causal_lens is None else causal_lens -+ ), -+ } -+ out, lse = cute_dsl_mla_decode( -+ return_lse=True, -+ **call_args, -+ **kwargs, -+ ) -+ return out, lse, split_kv -+ -+ -+def _assert_close_to_reference( -+ out: torch.Tensor, -+ lse: torch.Tensor, -+ ref_out: torch.Tensor, -+ ref_lse: torch.Tensor, -+ dtype: torch.dtype, -+) -> None: -+ if dtype == torch.float8_e4m3fn: -+ out_atol, out_rtol = 0.1, 0.1 -+ lse_atol, lse_rtol = 0.2, 0.1 -+ else: -+ out_atol = out_rtol = lse_atol = lse_rtol = 1e-2 -+ torch.testing.assert_close( -+ out.float(), ref_out, atol=out_atol, rtol=out_rtol -+ ) -+ torch.testing.assert_close( -+ lse.float(), ref_lse, atol=lse_atol, rtol=lse_rtol -+ ) -+ -+ -+def _assert_dcp_rank_merge_matches_reference( -+ query: torch.Tensor, -+ global_kv: torch.Tensor, -+ global_lens: torch.Tensor, -+ *, -+ cp_world: int, -+ dtype: torch.dtype, -+ is_var_seq: bool = False, -+) -> list[int]: -+ """Check every rank-local state and their final natural-log merge.""" -+ rank_outputs = [] -+ rank_lses = [] -+ split_kvs = [] -+ for cp_rank in range(cp_world): -+ out, lse, split_kv = _launch_rank( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=cp_world, -+ cp_rank=cp_rank, -+ enable_dcp=True, -+ is_var_seq=is_var_seq, -+ ) -+ ref_rank_out, ref_rank_lse = _reference_attention( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=cp_world, -+ cp_rank=cp_rank, -+ ) -+ _assert_close_to_reference( -+ out, lse, ref_rank_out, ref_rank_lse, dtype -+ ) -+ rank_outputs.append(out) -+ rank_lses.append(lse) -+ split_kvs.append(split_kv) -+ -+ merged_out, merged_lse = _merge_rank_outputs_natural_log( -+ rank_outputs, rank_lses -+ ) -+ ref_out, ref_lse = _reference_attention( -+ query, global_kv, global_lens -+ ) -+ _assert_close_to_reference( -+ merged_out, merged_lse, ref_out, ref_lse, dtype -+ ) -+ return split_kvs -+ -+ -+def test_dcp_flat_mask_matches_ceiling_divided_bound(): -+ """The division-free flattened predicate must match global coordinates.""" -+ for num_heads in (6, 12, 24, 48, 64, 96, 128): -+ for q_len in (1, 2, 4, 8): -+ for cp_world in (1, 2, 4): -+ for cp_rank in range(cp_world): -+ for global_len in (q_len, q_len + 1, 127, 128, 129): -+ local_len = _local_length( -+ global_len, cp_world, cp_rank -+ ) -+ for q_idx in range(q_len): -+ for local_key in range(local_len): -+ expected = ( -+ local_key * cp_world + cp_rank -+ < global_len - (q_len - 1) + q_idx -+ ) -+ for head in (0, num_heads - 1): -+ flat_q_row = q_idx * num_heads + head -+ assert ( -+ _flat_dcp_score_is_valid( -+ flat_q_row, -+ num_heads, -+ local_key, -+ q_len, -+ global_len, -+ cp_world, -+ cp_rank, -+ ) -+ == expected -+ ) -+ -+ -+def test_dcp_per_query_tile_dense_boundary_is_conservative(): -+ """A tile marked dense must be valid for every real row and local key.""" -+ k_tile = 128 -+ for num_heads in (6, 12, 24, 48, 64, 96, 128): -+ for q_len in (2, 4, 8, 32): -+ num_q_tiles = _ceil_div(q_len * num_heads, 128) -+ for cp_world in (1, 2, 4): -+ for cp_rank in range(cp_world): -+ global_len = max(q_len, 277) -+ local_len = _local_length( -+ global_len, cp_world, cp_rank -+ ) -+ global_positions = range( -+ cp_rank, global_len, cp_world -+ ) -+ for q_tile_idx in range(num_q_tiles): -+ first_flat_row = q_tile_idx * 128 -+ first_q = first_flat_row // num_heads -+ earliest_bound = bisect_left( -+ global_positions, -+ global_len - (q_len - 1) + first_q, -+ ) -+ effective_bound = min( -+ max(earliest_bound, 0), local_len -+ ) -+ first_mask_k_tile = effective_bound // k_tile -+ for k_tile_idx in range(first_mask_k_tile): -+ local_begin = k_tile_idx * k_tile -+ local_end = min(local_begin + k_tile, local_len) -+ for flat_q_row in range( -+ first_flat_row, -+ min( -+ first_flat_row + 128, -+ q_len * num_heads, -+ ), -+ ): -+ for local_key in range( -+ local_begin, local_end -+ ): -+ assert _flat_dcp_score_is_valid( -+ flat_q_row, -+ num_heads, -+ local_key, -+ q_len, -+ global_len, -+ cp_world, -+ cp_rank, -+ ) -+ -+def test_cute_dsl_mla_dcp_rejects_incomplete_static_contract(): -+ """Reject DCP calls that cannot produce mergeable rank-local states.""" -+ _skip_if_unsupported() -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ cute_dsl_mla_decode, -+ ) -+ -+ torch.manual_seed(40) -+ query, global_kv, global_lens = _make_inputs( -+ global_length=128, -+ q_len=4, -+ num_heads=96, -+ dtype=torch.bfloat16, -+ ) -+ call_args, _ = _prepare_rank_call( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ cp_rank=0, -+ ) -+ -+ ragged_call_args = {**call_args, "query": query[:, 0]} -+ with pytest.raises(ValueError, match="query must have shape"): -+ cute_dsl_mla_decode(**ragged_call_args) -+ with pytest.raises(ValueError, match="requires return_lse=True"): -+ cute_dsl_mla_decode( -+ **call_args, -+ enable_dcp=True, -+ cp_world=2, -+ cp_rank=0, -+ causal_seqlens_kv_global=global_lens, -+ ) -+ with pytest.raises( -+ ValueError, match="causal_seqlens_kv_global is required" -+ ): -+ cute_dsl_mla_decode( -+ **call_args, -+ enable_dcp=True, -+ cp_world=2, -+ cp_rank=0, -+ return_lse=True, -+ ) -+ with pytest.raises(TypeError, match="must be a torch.Tensor"): -+ cute_dsl_mla_decode( -+ **call_args, -+ enable_dcp=True, -+ cp_world=2, -+ cp_rank=0, -+ causal_seqlens_kv_global=[128], -+ return_lse=True, -+ ) -+ with pytest.raises(ValueError, match="0 <= cp_rank < cp_world"): -+ cute_dsl_mla_decode( -+ **call_args, -+ enable_dcp=True, -+ cp_world=2, -+ cp_rank=2, -+ causal_seqlens_kv_global=global_lens, -+ return_lse=True, -+ ) -+ with pytest.raises(ValueError, match="require enable_dcp=True"): -+ cute_dsl_mla_decode(**call_args, cp_world=2) -+ -+ -+def test_cute_dsl_mla_dcp_dispatch_is_strictly_monolithic(): -+ """DCP must never silently select a backend with non-DCP mask semantics.""" -+ from flashinfer.cute_dsl.attention.mla_dispatch import _resolve_impl -+ from flashinfer.trace.templates.attention import ( -+ trtllm_batch_decode_mla_trace_dispatch, -+ ) -+ -+ assert ( -+ _resolve_impl( -+ requested="auto", -+ kwargs={"enable_dcp": True, "cp_world": 2, "cp_rank": 0}, -+ ) -+ == "monolithic" -+ ) -+ with pytest.raises(ValueError, match="only supported by the monolithic"): -+ _resolve_impl( -+ requested="modular", -+ kwargs={"enable_dcp": True, "cp_world": 2, "cp_rank": 0}, -+ ) -+ with pytest.raises(ValueError, match="cannot be combined with 'sinks'"): -+ _resolve_impl( -+ requested="auto", -+ kwargs={ -+ "enable_dcp": True, -+ "cp_world": 2, -+ "cp_rank": 0, -+ "sinks": torch.empty(1), -+ }, -+ ) -+ with pytest.raises(NotImplementedError, match="does not yet represent"): -+ trtllm_batch_decode_mla_trace_dispatch(enable_dcp=True) -+ -+ -+def test_cute_dsl_mla_dcp_world1_matches_disabled(): -+ """World-one DCP must preserve the existing monolithic MLA result.""" -+ _skip_if_unsupported() -+ torch.manual_seed(41) -+ query, global_kv, global_lens = _make_inputs( -+ global_length=128, -+ q_len=4, -+ num_heads=96, -+ dtype=torch.bfloat16, -+ ) -+ baseline_out, baseline_lse, baseline_split = _launch_rank( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=1, -+ cp_rank=0, -+ enable_dcp=False, -+ ) -+ dcp_out, dcp_lse, dcp_split = _launch_rank( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=1, -+ cp_rank=0, -+ enable_dcp=True, -+ ) -+ assert baseline_split == dcp_split == 1 -+ torch.testing.assert_close(dcp_out, baseline_out, atol=0, rtol=0) -+ torch.testing.assert_close(dcp_lse, baseline_lse, atol=0, rtol=0) -+ -+ -+@pytest.mark.parametrize( -+ "dtype", -+ [torch.bfloat16, torch.float8_e4m3fn], -+ ids=["bf16", "fp8"], -+) -+def test_cute_dsl_mla_dcp_q1_h128_rank_merge(dtype): -+ """Cover the distinct Q1/H128 schedule for both kernel families.""" -+ _skip_if_unsupported() -+ torch.manual_seed(49) -+ query, global_kv, global_lens = _make_inputs( -+ global_length=129, -+ q_len=1, -+ num_heads=128, -+ dtype=dtype, -+ ) -+ split_kvs = _assert_dcp_rank_merge_matches_reference( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ dtype=dtype, -+ ) -+ assert split_kvs == [1, 1] -+ -+ -+@pytest.mark.parametrize("num_heads", [12, 24, 48]) -+def test_cute_dsl_mla_dcp_packed_head_counts(num_heads): -+ """Exercise DCP across the supported packed query-tile geometries.""" -+ _skip_if_unsupported() -+ torch.manual_seed(50 + num_heads) -+ query, global_kv, global_lens = _make_inputs( -+ global_length=130, -+ q_len=8, -+ num_heads=num_heads, -+ dtype=torch.bfloat16, -+ ) -+ split_kvs = _assert_dcp_rank_merge_matches_reference( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ dtype=torch.bfloat16, -+ ) -+ assert split_kvs == [1, 1] -+ -+ -+def test_cute_dsl_mla_dcp_public_api_autotune_profiles_causal_tensor(): -+ """The public runner must batch-sweep DCP metadata and return caller B1.""" -+ _skip_if_unsupported() -+ from flashinfer import autotune -+ from flashinfer.autotuner import AutoTuner -+ from flashinfer.mla._core import ( -+ trtllm_batch_decode_with_kv_cache_mla, -+ ) -+ -+ torch.manual_seed(47) -+ query, global_kv, global_lens = _make_inputs( -+ # Two local pages satisfy the common public dispatcher's aligned -+ # block-table contract for page_size=64. -+ global_length=256, -+ q_len=4, -+ num_heads=96, -+ dtype=torch.bfloat16, -+ ) -+ call_args, split_kv = _prepare_rank_call( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ cp_rank=0, -+ ) -+ assert split_kv == 1 -+ -+ public_args = { -+ "query": query, -+ "kv_cache": call_args["kv_cache"], -+ "workspace_buffer": call_args["workspace_buffer"], -+ "qk_nope_head_dim": 128, -+ "kv_lora_rank": _LATENT_DIM, -+ "qk_rope_head_dim": _ROPE_DIM, -+ "block_tables": call_args["block_tables"], -+ "seq_lens": call_args["seq_lens"], -+ "max_seq_len": call_args["max_seq_len"], -+ "bmm1_scale": call_args["softmax_scale"], -+ "bmm2_scale": 1.0, -+ "backend": "auto", -+ "cute_dsl_impl": "monolithic", -+ "is_var_seq": False, -+ "enable_pdl": False, -+ "return_lse": True, -+ "enable_dcp": True, -+ "cp_world": 2, -+ "cp_rank": 0, -+ } -+ with pytest.raises(TypeError, match="must be a torch.Tensor"): -+ trtllm_batch_decode_with_kv_cache_mla( -+ **public_args, -+ causal_seqlens_kv_global=[256], -+ ) -+ -+ AutoTuner.get().clear_cache() -+ try: -+ with mock.patch( -+ "flashinfer.mla._core._compute_mla_decode_buckets", -+ return_value=(2,), -+ ), autotune(tune_mode=True): -+ out, lse = trtllm_batch_decode_with_kv_cache_mla( -+ **public_args, -+ causal_seqlens_kv_global=global_lens, -+ ) -+ finally: -+ AutoTuner.get().clear_cache() -+ -+ ref_out, ref_lse = _reference_attention( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ cp_rank=0, -+ ) -+ _assert_close_to_reference( -+ out, -+ lse.view_as(ref_lse), -+ ref_out, -+ ref_lse, -+ torch.bfloat16, -+ ) -+ -+ -+@pytest.mark.parametrize( -+ "dtype", -+ [torch.bfloat16, torch.float8_e4m3fn], -+ ids=["bf16", "fp8"], -+) -+def test_cute_dsl_mla_dcp_rank_mask_and_merge(dtype): -+ """Cover the rank-sensitive G128/Q4/H96 boundary and rank merge.""" -+ _skip_if_unsupported() -+ torch.manual_seed(42) -+ query, global_kv, global_lens = _make_inputs( -+ global_length=128, -+ q_len=4, -+ num_heads=96, -+ dtype=dtype, -+ ) -+ split_kvs = _assert_dcp_rank_merge_matches_reference( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ dtype=dtype, -+ ) -+ assert split_kvs == [1, 1] -+ -+ # In particular, query zero on rank one has local bound 62, not 63. -+ assert ( -+ _local_causal_bound(128, 4, 0, cp_world=2, cp_rank=1) -+ == 62 -+ ) -+ -+ -+@pytest.mark.parametrize( -+ "dtype", -+ [torch.float16, torch.float8_e4m3fn], -+ ids=["fp16", "fp8"], -+) -+def test_cute_dsl_mla_dcp_world4_variable_batch_h6(dtype): -+ """Cover W4, heterogeneous local tails, FP16, and packed H6 query rows.""" -+ _skip_if_unsupported() -+ torch.manual_seed(46) -+ query, global_kv, global_lens = _make_batched_inputs( -+ global_lengths=(65, 130, 259), -+ q_len=8, -+ num_heads=6, -+ dtype=dtype, -+ ) -+ split_kvs = _assert_dcp_rank_merge_matches_reference( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=4, -+ dtype=dtype, -+ is_var_seq=True, -+ ) -+ assert split_kvs == [1] * 4 -+ -+ -+@pytest.mark.parametrize( -+ "cp_world,global_length,dtype", -+ [ -+ pytest.param(4, 4 * 1024 + 3, torch.bfloat16, id="w4-k4k-bf16"), -+ pytest.param( -+ 4, -+ 4 * 1024 + 3, -+ torch.float8_e4m3fn, -+ id="w4-k4k-fp8", -+ ), -+ pytest.param(8, 8 * 1024 + 5, torch.bfloat16, id="w8-k8k-bf16"), -+ pytest.param( -+ 8, -+ 8 * 1024 + 5, -+ torch.float8_e4m3fn, -+ id="w8-k8k-fp8", -+ ), -+ pytest.param(16, 16 * 1024 + 9, torch.bfloat16, id="w16-k16k-bf16"), -+ pytest.param( -+ 16, -+ 16 * 1024 + 9, -+ torch.float8_e4m3fn, -+ id="w16-k16k-fp8", -+ ), -+ ], -+) -+def test_cute_dsl_mla_dcp_rank_scale_merge( -+ cp_world, global_length, dtype -+): -+ """Cover W4/W8/W16 four-query speculative decode and uneven split-KV.""" -+ _skip_if_unsupported() -+ torch.manual_seed(51 + cp_world) -+ query, global_kv, global_lens = _make_inputs( -+ global_length=global_length, -+ q_len=4, -+ num_heads=96, -+ dtype=dtype, -+ ) -+ -+ base_local_length, tail_ranks = divmod(global_length, cp_world) -+ for cp_rank in range(cp_world): -+ expected_local_length = base_local_length + ( -+ 1 if cp_rank < tail_ranks else 0 -+ ) -+ assert ( -+ _local_length(global_length, cp_world, cp_rank) -+ == expected_local_length -+ ) -+ -+ split_kvs = _assert_dcp_rank_merge_matches_reference( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=cp_world, -+ dtype=dtype, -+ ) -+ assert all(split_kv > 1 for split_kv in split_kvs) -+ -+ -+def test_cute_dsl_mla_dcp_cuda_graph_reads_updated_causal_bound(): -+ """A captured launch must read causal bounds mutated in place on replay.""" -+ _skip_if_unsupported() -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ cute_dsl_mla_decode, -+ ) -+ -+ torch.manual_seed(45) -+ query, global_kv, global_lens = _make_inputs( -+ global_length=128, -+ q_len=4, -+ num_heads=96, -+ dtype=torch.bfloat16, -+ ) -+ call_args, split_kv = _prepare_rank_call( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ cp_rank=1, -+ ) -+ assert split_kv == 1 -+ out = torch.empty( -+ 1, -+ 4, -+ 96, -+ _LATENT_DIM, -+ dtype=torch.bfloat16, -+ device=query.device, -+ ) -+ lse = torch.empty( -+ 1, 4, 96, dtype=torch.float32, device=query.device -+ ) -+ dcp_args = { -+ "enable_dcp": True, -+ "cp_world": 2, -+ "cp_rank": 1, -+ "causal_seqlens_kv_global": global_lens, -+ "return_lse": True, -+ "out": out, -+ "lse": lse, -+ } -+ -+ # Compile and initialize all persistent buffers before capture. -+ returned_out, returned_lse = cute_dsl_mla_decode( -+ **call_args, **dcp_args -+ ) -+ assert returned_out is out -+ assert returned_lse is lse -+ torch.cuda.synchronize() -+ -+ graph = torch.cuda.CUDAGraph() -+ with torch.cuda.graph(graph): -+ cute_dsl_mla_decode(**call_args, **dcp_args) -+ lse_at_128 = lse.clone() -+ -+ # Keep the physical local cache and seq_lens fixed at the G=128 shard. -+ # The G=126 replay must mask its now-extra global keys 126 and 127. -+ global_lens.fill_(126) -+ graph.replay() -+ torch.cuda.synchronize() -+ assert not torch.equal(lse, lse_at_128) -+ ref_out, ref_lse = _reference_attention( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ cp_rank=1, -+ ) -+ _assert_close_to_reference( -+ out, lse, ref_out, ref_lse, torch.bfloat16 -+ ) -+ -+ -+@pytest.mark.parametrize( -+ "dtype", -+ [torch.bfloat16, torch.float8_e4m3fn], -+ ids=["bf16", "fp8"], -+) -+def test_cute_dsl_mla_dcp_empty_rank_row(dtype): -+ """A rank with no visible key must contribute O=0 and LSE=-inf.""" -+ _skip_if_unsupported() -+ torch.manual_seed(43) -+ query, global_kv, global_lens = _make_inputs( -+ global_length=4, -+ q_len=4, -+ num_heads=96, -+ dtype=dtype, -+ ) -+ out, lse, split_kv = _launch_rank( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ cp_rank=1, -+ enable_dcp=True, -+ ) -+ assert split_kv == 1 -+ assert torch.equal(out[:, 0], torch.zeros_like(out[:, 0])) -+ assert torch.isneginf(lse[:, 0]).all() -+ ref_out, ref_lse = _reference_attention( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ cp_rank=1, -+ ) -+ _assert_close_to_reference( -+ out, lse, ref_out, ref_lse, dtype -+ ) -+ -+ # Also cover a physically empty cyclic shard and the cross-rank merge for -+ # early rows that have no visible key on any rank. The empty shard still -+ # owns one padding page, but seq_lens=0, so it may not load or reduce K. -+ query, global_kv, global_lens = _make_inputs( -+ global_length=1, -+ q_len=4, -+ num_heads=96, -+ dtype=dtype, -+ ) -+ split_kvs = _assert_dcp_rank_merge_matches_reference( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ dtype=dtype, -+ ) -+ assert split_kvs == [1, 1] -+ -+ -+@pytest.mark.parametrize( -+ "dtype", -+ [torch.bfloat16, torch.float8_e4m3fn], -+ ids=["bf16", "fp8"], -+) -+def test_cute_dsl_mla_dcp_split_kv_rank_merge(dtype): -+ """DCP must preserve natural-log LSE through standalone split reduction.""" -+ _skip_if_unsupported() -+ torch.manual_seed(44) -+ query, global_kv, global_lens = _make_inputs( -+ global_length=4096, -+ q_len=4, -+ num_heads=96, -+ dtype=dtype, -+ ) -+ split_kvs = _assert_dcp_rank_merge_matches_reference( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ dtype=dtype, -+ ) -+ assert all(split_kv > 1 for split_kv in split_kvs) -+ -+ # Preserve the 2048-token physical shard and its split geometry, but make -+ # every rank-1 key causally invisible. This exercises the all-empty split -+ # reducer rather than the direct zero-tile output path above. -+ empty_causal_lens = torch.tensor( -+ [1], dtype=torch.int32, device=query.device -+ ) -+ empty_out, empty_lse, empty_split_kv = _launch_rank( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ cp_rank=1, -+ enable_dcp=True, -+ causal_lens=empty_causal_lens, -+ ) -+ assert empty_split_kv == split_kvs[1] > 1 -+ assert torch.equal(empty_out, torch.zeros_like(empty_out)) -+ assert torch.isneginf(empty_lse).all() -+ ref_empty_out, ref_empty_lse = _reference_attention( -+ query, -+ global_kv, -+ empty_causal_lens, -+ cp_world=2, -+ cp_rank=1, -+ ) -+ _assert_close_to_reference( -+ empty_out, -+ empty_lse, -+ ref_empty_out, -+ ref_empty_lse, -+ dtype, -+ ) -+ -+ -+def test_cute_dsl_mla_dcp_pdl_all_empty_split_reduction(): -+ """PDL must order the all-empty producer path before split reduction.""" -+ _skip_if_unsupported() -+ from flashinfer.utils import device_support_pdl -+ -+ if not device_support_pdl(torch.device("cuda")): -+ pytest.skip("Programmatic dependent launch is not supported") -+ -+ torch.manual_seed(48) -+ query, global_kv, global_lens = _make_inputs( -+ global_length=4096, -+ q_len=4, -+ num_heads=96, -+ dtype=torch.bfloat16, -+ ) -+ empty_causal_lens = torch.tensor( -+ [1], dtype=torch.int32, device=query.device -+ ) -+ out, lse, split_kv = _launch_rank( -+ query, -+ global_kv, -+ global_lens, -+ cp_world=2, -+ cp_rank=1, -+ enable_dcp=True, -+ causal_lens=empty_causal_lens, -+ enable_pdl=True, -+ ) -+ assert split_kv > 1 -+ assert torch.equal(out, torch.zeros_like(out)) -+ assert torch.isneginf(lse).all() -diff --git a/tests/attention/test_cute_dsl_mla_decode.py b/tests/attention/test_cute_dsl_mla_decode.py -index e8549a72..f52dee4f 100644 ---- a/tests/attention/test_cute_dsl_mla_decode.py -+++ b/tests/attention/test_cute_dsl_mla_decode.py -@@ -83,15 +83,19 @@ def torch_reference_mla( - - outputs = [] - lses = [] -+ # Copy the small metadata arrays once. Calling CUDA ``.item()`` for every -+ # request and page serializes hundreds of device synchronizations at B128. -+ cache_seqs_host = cache_seqs.tolist() -+ page_table_host = page_table.tolist() - for b in range(B): -- seq_len = cache_seqs[b].item() -+ seq_len = cache_seqs_host[b] - num_pages_needed = (seq_len + page_size - 1) // page_size - - # Gather KV for this batch via page table -- page_indices = page_table[b, :num_pages_needed] -+ page_indices = page_table_host[b][:num_pages_needed] - kv_indices = [] - for p in page_indices: -- start = p.item() * page_size -+ start = p * page_size - kv_indices.extend(range(start, start + page_size)) - kv_indices = kv_indices[:seq_len] - kv_indices_t = torch.tensor(kv_indices, device=q_nope.device) -@@ -254,28 +258,21 @@ def test_cute_dsl_mla_decode_fp16( - torch.testing.assert_close(lse, ref_lse, atol=1e-2, rtol=1e-2) - - --# Exercises the spec-decoding (MTP) causal mask + fold_sq path: num_heads < 128 --# forces the kernel to pack F = compute_fold_sq_ratio(H, q_len, 128) tokens of --# q_len into the head dim so the 128-wide MMA-M tile is fully populated. --# (H=128, q_len=any) → F=1 (no fold), (H=64, q_len=2) → F=2, (H=64, q_len=4) → F=2, --# (H=32, q_len=4) → F=4, (H=32, q_len=2) → F=2. All paths share the same --# kernel; the MTP causal mask is applied uniformly for q_len > 1. --# Monolithic-only: the modular path doesn't implement fold_sq or the MTP mask. --@pytest.mark.parametrize("batch_size", [1, 4]) --@pytest.mark.parametrize("seq_len_k", [128, 1024]) --@pytest.mark.parametrize("num_heads", [16, 32, 64]) --@pytest.mark.parametrize("q_len", [2, 4]) --@pytest.mark.parametrize("dtype", [torch.float16, torch.float8_e4m3fn]) --def test_cute_dsl_mla_decode_fold_sq( -- batch_size, seq_len_k, num_heads, q_len, dtype, cute_dsl_impl -+def _run_padded_q_tile_case( -+ batch_size, -+ seq_len_k, -+ num_heads, -+ q_len, -+ dtype, -+ is_var_seq=False, -+ enable_pdl=None, -+ out_dtype=None, -+ query_token_stride=1, -+ via_public_api=False, - ): -- """Verify the MTP causal mask + fold_sq packing for H ≤ 128 and q_len > 1.""" -- if cute_dsl_impl != "monolithic": -- pytest.skip("fold_sq / MTP causal mask are monolithic-only features") -+ """Run and reference-check one monolithic packed-query configuration.""" - skip_if_unsupported() - -- from flashinfer.cute_dsl.attention import cute_dsl_mla_decode -- - torch.manual_seed(42) - device = torch.device("cuda") - -@@ -288,17 +285,19 @@ def test_cute_dsl_mla_decode_fold_sq( - - # torch.randn doesn't support fp8; for FP8 inputs create as fp16 then convert. - is_fp8 = dtype == torch.float8_e4m3fn -+ query_storage = torch.randn( -+ batch_size, -+ q_len * query_token_stride, -+ num_heads, -+ D_qk, -+ dtype=torch.float16 if is_fp8 else dtype, -+ device=device, -+ ) - if is_fp8: -- query = ( -- torch.randn( -- batch_size, q_len, num_heads, D_qk, dtype=torch.float16, device=device -- ) -- * 0.1 -- ).to(torch.float8_e4m3fn) -- else: -- query = torch.randn( -- batch_size, q_len, num_heads, D_qk, dtype=dtype, device=device -- ) -+ query_storage = (query_storage * 0.1).to(torch.float8_e4m3fn) -+ query = query_storage[:, ::query_token_stride] -+ if query_token_stride != 1: -+ assert query.stride(1) != num_heads * query.stride(2) - - num_pages_per_batch = (seq_len_k + page_size - 1) // page_size - total_pages = num_pages_per_batch * batch_size + 10 -@@ -319,25 +318,79 @@ def test_cute_dsl_mla_decode_fold_sq( - for p in range(num_pages_per_batch): - block_tables[b, p] = b * num_pages_per_batch + p - -- seq_lens = torch.full((batch_size,), seq_len_k, dtype=torch.int32, device=device) -- -- workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8, device=device) -+ if is_var_seq: -+ seq_lens = torch.tensor( -+ [max(page_size, seq_len_k - b * 37) for b in range(batch_size)], -+ dtype=torch.int32, -+ device=device, -+ ) -+ else: -+ seq_lens = torch.full( -+ (batch_size,), seq_len_k, dtype=torch.int32, device=device -+ ) - -- out = cute_dsl_mla_decode( -- query=query, -- kv_cache=kv_cache, -- workspace_buffer=workspace_buffer, -- kv_lora_rank=latent_dim, -- qk_rope_head_dim=rope_dim, -- block_tables=block_tables, -- seq_lens=seq_lens, -- max_seq_len=seq_len_k, -- softmax_scale=softmax_scale, -- output_scale=output_scale, -- is_var_seq=False, -- cute_dsl_impl=cute_dsl_impl, -+ workspace_factory = torch.zeros if via_public_api else torch.empty -+ workspace_buffer = workspace_factory( -+ 256 * 1024 * 1024, dtype=torch.int8, device=device - ) - -+ if via_public_api: -+ assert out_dtype is None -+ from flashinfer.mla import trtllm_batch_decode_with_kv_cache_mla -+ -+ lse_out = torch.empty( -+ batch_size, q_len, num_heads, dtype=torch.float32, device=device -+ ) -+ out, lse = trtllm_batch_decode_with_kv_cache_mla( -+ query=query, -+ kv_cache=kv_cache, -+ workspace_buffer=workspace_buffer, -+ qk_nope_head_dim=latent_dim, -+ kv_lora_rank=latent_dim, -+ qk_rope_head_dim=rope_dim, -+ block_tables=block_tables, -+ seq_lens=seq_lens, -+ max_seq_len=seq_len_k, -+ bmm1_scale=softmax_scale, -+ bmm2_scale=output_scale, -+ backend="cute-dsl", -+ is_var_seq=is_var_seq, -+ enable_pdl=enable_pdl, -+ lse=lse_out, -+ return_lse=True, -+ cute_dsl_impl="monolithic", -+ ) -+ else: -+ from flashinfer.cute_dsl.attention import cute_dsl_mla_decode -+ -+ out, lse = cute_dsl_mla_decode( -+ query=query, -+ kv_cache=kv_cache, -+ workspace_buffer=workspace_buffer, -+ kv_lora_rank=latent_dim, -+ qk_rope_head_dim=rope_dim, -+ block_tables=block_tables, -+ seq_lens=seq_lens, -+ max_seq_len=seq_len_k, -+ softmax_scale=softmax_scale, -+ output_scale=output_scale, -+ is_var_seq=is_var_seq, -+ enable_pdl=enable_pdl, -+ cute_dsl_impl="monolithic", -+ return_lse=True, -+ out_dtype=out_dtype, -+ ) -+ -+ if via_public_api: -+ expected_out_dtype = torch.bfloat16 -+ elif out_dtype is not None: -+ expected_out_dtype = out_dtype -+ elif is_fp8: -+ expected_out_dtype = torch.bfloat16 -+ else: -+ expected_out_dtype = dtype -+ assert out.dtype == expected_out_dtype -+ - # FP8 input → BF16 output (default), so do the reference in FP32 with wider tolerance. - if is_fp8: - kv_flat = kv_cache.reshape(-1, D_qk).to(torch.float32) -@@ -351,7 +404,7 @@ def test_cute_dsl_mla_decode_fold_sq( - c_rope_ref = kv_flat[:, latent_dim:] - - # Monolithic-only test — kernel always applies the MTP causal mask here. -- ref_out = torch_reference_mla( -+ ref_out, ref_lse = torch_reference_mla( - q_nope, - q_rope, - c_latent_ref, -@@ -362,6 +415,7 @@ def test_cute_dsl_mla_decode_fold_sq( - output_scale, - page_size, - apply_mtp_mask=True, -+ return_lse=True, - ) - - if is_fp8: -@@ -369,45 +423,518 @@ def test_cute_dsl_mla_decode_fold_sq( - torch.testing.assert_close( - out.to(torch.float32), ref_out.to(torch.float32), atol=0.1, rtol=0.1 - ) -+ torch.testing.assert_close(lse, ref_lse, atol=0.2, rtol=0.1) - else: -- ref_out_cast = ref_out.to(dtype) -+ ref_out_cast = ref_out.to(out.dtype) - torch.testing.assert_close(out, ref_out_cast, atol=1e-2, rtol=1e-2) -+ torch.testing.assert_close(lse, ref_lse, atol=1e-2, rtol=1e-2) -+ -+ -+# Exercises the spec-decoding (MTP) causal mask + flat query-row packing. -+# A cooperative 2-CTA tile owns 128 consecutive (token, head) rows and may -+# cross token boundaries; only the final flattened tile is safely padded. -+# All paths share the same kernel, and the MTP mask applies for q_len > 1. -+# Monolithic-only: the modular path doesn't implement packed query tiles or MTP. -+@pytest.mark.parametrize( -+ "num_heads,q_len,dtype", -+ [ -+ pytest.param(64, 3, torch.bfloat16, id="h64-sq3-bf16"), -+ pytest.param(24, 13, torch.float8_e4m3fn, id="h24-sq13-fp8"), -+ pytest.param(12, 11, torch.bfloat16, id="h12-sq11-bf16-tail4"), -+ ], -+) -+def test_cute_dsl_mla_decode_padded_q_tiles(num_heads, q_len, dtype): -+ """Cover cross-token rows and final tails in both kernel families.""" -+ _run_padded_q_tile_case(1, 1024, num_heads, q_len, dtype) -+ -+ -+@pytest.mark.parametrize( -+ "dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"] -+) -+def test_cute_dsl_mla_decode_per_q_tile_mask_boundary(dtype): -+ """Use each M128 tile's first token when selecting the dense mask path.""" -+ skip_if_unsupported() -+ -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ _get_split_kv_and_workspace_size, -+ cute_dsl_mla_decode, -+ ) -+ from flashinfer.cute_dsl.utils import get_num_sm -+ -+ batch_size, q_len, num_heads = 1, 32, 6 -+ seq_len_k, page_size = 144, 64 -+ latent_dim, rope_dim = 512, 64 -+ d_qk = latent_dim + rope_dim -+ device = torch.device("cuda") -+ -+ # H6/Sq32 has two M128 tiles. Token 21 straddles their boundary: heads -+ # 0-1 are in tile 0 and heads 2-5 are in tile 1. Tile 1 can treat keys -+ # 0:128 as dense from its first token (21), while tile 0 must still mask -+ # keys 113:128 for its earlier tokens. Using tile 0's last token for the -+ # coarse decision would expose those marked values incorrectly. -+ query = torch.zeros(batch_size, q_len, num_heads, d_qk, dtype=dtype, device=device) -+ num_pages = (seq_len_k + page_size - 1) // page_size -+ kv_cache = torch.zeros(num_pages, page_size, d_qk, dtype=dtype, device=device) -+ kv_cache.view(-1, d_qk)[113:128, 0] = 4.0 -+ block_tables = torch.arange(num_pages, dtype=torch.int32, device=device).view(1, -1) -+ seq_lens = torch.full((batch_size,), seq_len_k, dtype=torch.int32, device=device) -+ -+ split_kv, workspace_size = _get_split_kv_and_workspace_size( -+ batch_size, -+ q_len, -+ num_heads, -+ latent_dim, -+ get_num_sm(device), -+ seq_len_k, -+ ) -+ assert split_kv == 2 -+ workspace = torch.empty(workspace_size, dtype=torch.int8, device=device) -+ -+ out, lse = cute_dsl_mla_decode( -+ query=query, -+ kv_cache=kv_cache, -+ workspace_buffer=workspace, -+ kv_lora_rank=latent_dim, -+ qk_rope_head_dim=rope_dim, -+ block_tables=block_tables, -+ seq_lens=seq_lens, -+ max_seq_len=seq_len_k, -+ softmax_scale=1.0 / (latent_dim**0.5), -+ return_lse=True, -+ enable_pdl=False, -+ ) -+ -+ bounds = torch.arange(q_len, dtype=torch.float32, device=device) + ( -+ seq_len_k - q_len + 1 -+ ) -+ marked_visible = (bounds - 113).clamp(min=0, max=15) -+ expected_out = torch.zeros_like(out, dtype=torch.float32) -+ expected_out[..., 0] = (4.0 * marked_visible / bounds).view(1, q_len, 1) -+ expected_lse = torch.log(bounds).view(1, q_len, 1).expand_as(lse) -+ -+ torch.testing.assert_close(out.float(), expected_out, atol=5e-3, rtol=5e-3) -+ torch.testing.assert_close(lse, expected_lse, atol=2e-3, rtol=2e-3) -+ -+ -+@pytest.mark.parametrize("num_heads", [12, 24, 48, 96]) -+@pytest.mark.parametrize("q_len", [1, 2, 4, 8]) -+@pytest.mark.parametrize( -+ "dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"] -+) -+def test_cute_dsl_mla_decode_packed_q_accuracy_matrix(num_heads, q_len, dtype): -+ """Qualify the supported packed-query, causal-mask, and auto-split matrix.""" -+ # Keep the batch sweep within one pytest item while all 128 requested -+ # Cartesian cells are reference checked. Split-KV and query packing are -+ # selected automatically by the public launcher. -+ for batch_size in (1, 4, 16, 128): -+ _run_padded_q_tile_case( -+ batch_size=batch_size, -+ seq_len_k=1024, -+ num_heads=num_heads, -+ q_len=q_len, -+ dtype=dtype, -+ ) -+ -+ -+def test_cute_dsl_mla_decode_fp8_persistent_multi_work_boundaries(): -+ """Reuse the FP8 two-softmax pipelines across persistent work items.""" -+ # B128 exceeds one resident wave of 2CTA clusters on supported Blackwell -+ # devices. These cases exercise one-, even-, and odd-K-tile work items; -+ # the packed-Q matrix above additionally covers the eight-tile boundary. -+ for seq_len_k, enable_pdl in ((128, True), (256, False), (384, True)): -+ _run_padded_q_tile_case( -+ batch_size=128, -+ seq_len_k=seq_len_k, -+ num_heads=96, -+ q_len=1, -+ dtype=torch.float8_e4m3fn, -+ enable_pdl=enable_pdl, -+ ) -+ -+ -+def test_cute_dsl_mla_decode_fp8_variable_seq_order_boundaries(): -+ """Balance both FP8 softmax groups for mixed nonpersistent K parity.""" -+ _run_padded_q_tile_case( -+ batch_size=10, -+ seq_len_k=385, -+ num_heads=96, -+ q_len=8, -+ dtype=torch.float8_e4m3fn, -+ is_var_seq=True, -+ enable_pdl=False, -+ ) -+ -+ -+@pytest.mark.parametrize( -+ "dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"] -+) -+def test_cute_dsl_mla_decode_h96_max_split_reducer_capacity(dtype): -+ """Reference-check output and LSE at the static reducer's 32-split cap.""" -+ _run_padded_q_tile_case( -+ batch_size=1, -+ seq_len_k=8192, -+ num_heads=96, -+ q_len=1, -+ dtype=dtype, -+ ) -+ -+ -+def test_cute_dsl_mla_decode_h96_sq8_nonempty_split_reducer(): -+ """Reference-check the normalized H96/Sq8 long-K split-reducer path.""" -+ _run_padded_q_tile_case( -+ batch_size=1, -+ seq_len_k=8192, -+ num_heads=96, -+ q_len=8, -+ dtype=torch.bfloat16, -+ ) -+ -+ -+@pytest.mark.parametrize( -+ "dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"] -+) -+def test_cute_dsl_mla_decode_h96_odd_split_reducer_pdl_off(dtype): -+ """Cover adaptive D4 with an odd 17-split prefix and PDL disabled.""" -+ _run_padded_q_tile_case( -+ batch_size=1, -+ seq_len_k=4097, -+ num_heads=96, -+ q_len=1, -+ dtype=dtype, -+ enable_pdl=False, -+ ) -+ -+ -+def test_cute_dsl_mla_decode_variable_seq_d2_reducer(): -+ """Cover adaptive D2 when batches have distinct non-power split prefixes.""" -+ _run_padded_q_tile_case( -+ batch_size=2, -+ seq_len_k=4097, -+ num_heads=24, -+ q_len=1, -+ dtype=torch.bfloat16, -+ is_var_seq=True, -+ enable_pdl=False, -+ ) -+ -+ -+def test_cute_dsl_mla_decode_padded_q_tile_direct_output(): -+ """Exercise H48/Sq3 full/tail padding with split_kv=1 (no reducer).""" -+ skip_if_unsupported() -+ -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ _get_split_kv_and_workspace_size, -+ ) -+ from flashinfer.cute_dsl.utils import get_num_sm -+ -+ num_q_tiles = 2 -+ num_sm = get_num_sm(torch.device("cuda")) -+ batch_size = (num_sm + 2 * num_q_tiles - 1) // (2 * num_q_tiles) -+ split_kv, workspace_size = _get_split_kv_and_workspace_size( -+ batch_size, 3, 48, 512, num_sm -+ ) -+ assert split_kv == 1 -+ assert workspace_size == 0 -+ -+ _run_padded_q_tile_case( -+ batch_size=batch_size, -+ seq_len_k=128, -+ num_heads=48, -+ q_len=3, -+ dtype=torch.float16, -+ ) -+ -+ -+def test_cute_dsl_mla_decode_padded_q_tile_variable_seq(): -+ """Cover nonpersistent variable-K scheduling with a padded H24/Sq6 tail.""" -+ _run_padded_q_tile_case( -+ batch_size=3, -+ seq_len_k=385, -+ num_heads=24, -+ q_len=6, -+ dtype=torch.float16, -+ is_var_seq=True, -+ enable_pdl=False, -+ ) - - --def test_compute_fold_sq_ratio(): -- """Unit test the static helper used by both run() and the wrapper.""" -+def test_cute_dsl_mla_decode_padded_q_tile_fp8_output(): -+ """Qualify packed-query FP8 input and FP8 output together.""" -+ _run_padded_q_tile_case( -+ batch_size=1, -+ seq_len_k=128, -+ num_heads=64, -+ q_len=3, -+ dtype=torch.float8_e4m3fn, -+ out_dtype=torch.float8_e4m3fn, -+ ) -+ -+ -+def test_cute_dsl_mla_decode_padded_q_tile_strided_query(): -+ """Cover the contiguous fallback for a token-strided query view.""" -+ _run_padded_q_tile_case( -+ batch_size=1, -+ seq_len_k=128, -+ num_heads=48, -+ q_len=3, -+ dtype=torch.float16, -+ query_token_stride=2, -+ ) -+ -+ -+@pytest.mark.parametrize("buffer_name", ["out", "lse"]) -+def test_cute_dsl_mla_decode_rejects_token_gapped_output(buffer_name): -+ """Reject output views that cannot represent flat rows across tokens.""" - if not is_cute_dsl_available(): - pytest.skip("CuTe DSL not available") -- from flashinfer.cute_dsl.attention.monolithic.mla_decode_fp16 import ( -- BlackwellMultiHeadLatentAttentionForwardFP16 as FP16, -+ -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ cute_dsl_mla_decode, - ) -- from flashinfer.cute_dsl.attention.monolithic.mla_decode_fp8 import ( -- BlackwellMultiHeadLatentAttentionForwardFP8 as FP8, -+ -+ batch_size, q_len, num_heads = 1, 2, 96 -+ latent_dim, rope_dim = 512, 64 -+ query = torch.empty( -+ batch_size, -+ q_len, -+ num_heads, -+ latent_dim + rope_dim, -+ dtype=torch.bfloat16, -+ ) -+ out = torch.empty(batch_size, q_len * 2, num_heads, latent_dim)[:, ::2] -+ lse = torch.empty(batch_size, q_len * 2, num_heads, dtype=torch.float32)[:, ::2] -+ kwargs = {buffer_name: out if buffer_name == "out" else lse} -+ -+ with pytest.raises(ValueError, match=rf"{buffer_name} must be contiguous"): -+ cute_dsl_mla_decode( -+ query=query, -+ kv_cache=torch.empty(1, 64, latent_dim + rope_dim, dtype=torch.bfloat16), -+ workspace_buffer=torch.empty(0, dtype=torch.int8), -+ kv_lora_rank=latent_dim, -+ qk_rope_head_dim=rope_dim, -+ block_tables=torch.zeros(1, 1, dtype=torch.int32), -+ seq_lens=torch.tensor([64], dtype=torch.int32), -+ max_seq_len=64, -+ softmax_scale=1.0, -+ **kwargs, -+ ) -+ -+ -+def test_cute_dsl_mla_decode_padded_q_tile_via_public_api(): -+ """Exercise the reported H96/Sq8 shape through the public dispatcher.""" -+ _run_padded_q_tile_case( -+ batch_size=1, -+ seq_len_k=128, -+ num_heads=96, -+ q_len=8, -+ dtype=torch.bfloat16, -+ enable_pdl=False, -+ via_public_api=True, -+ ) -+ -+ -+def test_compute_q_tile_layout(): -+ """Unit test the shared host/kernel query-tile geometry.""" -+ if not is_cute_dsl_available(): -+ pytest.skip("CuTe DSL not available") -+ from flashinfer.cute_dsl.attention.monolithic.mla_helpers import ( -+ compute_q_tile_layout, - ) - - cases = [ -- # (num_heads, seq_len_q, m_tile, expected) -- (128, 1, 128, 1), # H == M_tile → no fold -- (128, 4, 128, 1), # H == M_tile → no fold -- (64, 1, 128, 1), # seq_len_q=1 → F=1 -- (64, 2, 128, 2), # exact divisor, H*F=128 ≤ M_tile -- (64, 4, 128, 2), # H*F ≤ 128 caps F at 2; 4 % 2 == 0 -- (64, 3, 128, 1), # 3's only divisors are 1 and 3; H*3=192 > M_tile → F=1 -- (32, 4, 128, 4), # tighter pack: F=4, H*F=128 -- (32, 8, 128, 4), # capped by M_tile/H = 4 -- (32, 3, 128, 3), # max_fold=min(3, 4)=3; 3 % 3 == 0 → F=3 -- (16, 8, 128, 8), # max_fold=min(8, 8)=8; 8 % 8 == 0 → F=8 -- (16, 6, 128, 6), # max_fold=min(6, 8)=6; 6 % 6 == 0 → F=6 -+ # (H, Sq, M, (total_rows, num_tiles, tail_rows)) -+ (128, 3, 128, (384, 3, 128)), -+ (96, 1, 128, (96, 1, 96)), -+ (96, 3, 128, (288, 3, 32)), -+ (48, 8, 128, (384, 3, 128)), -+ (12, 11, 128, (132, 2, 4)), - ] - for H, S_q, m_tile, expected in cases: -- assert FP16.compute_fold_sq_ratio(H, S_q, m_tile) == expected, ( -- f"FP16.compute_fold_sq_ratio({H}, {S_q}, {m_tile}) " -- f"= {FP16.compute_fold_sq_ratio(H, S_q, m_tile)}, expected {expected}" -+ assert compute_q_tile_layout(H, S_q, m_tile) == expected, ( -+ f"compute_q_tile_layout({H}, {S_q}, {m_tile}) " -+ f"= {compute_q_tile_layout(H, S_q, m_tile)}, expected {expected}" - ) -- assert FP8.compute_fold_sq_ratio(H, S_q, m_tile) == expected, ( -- f"FP8.compute_fold_sq_ratio({H}, {S_q}, {m_tile}) " -- f"= {FP8.compute_fold_sq_ratio(H, S_q, m_tile)}, expected {expected}" -+ -+ for invalid in [(0, 1, 128), (129, 1, 128), (64, 0, 128), (64, 1, 0)]: -+ with pytest.raises(ValueError): -+ compute_q_tile_layout(*invalid) -+ -+ -+def test_nonpersistent_grid_y_limit(): -+ """Reject only nonpersistent grids beyond CUDA's Y-dimension limit.""" -+ if not is_cute_dsl_available(): -+ pytest.skip("CuTe DSL not available") -+ -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ _validate_nonpersistent_grid_y, -+ ) -+ -+ _validate_nonpersistent_grid_y(8_191, 8, is_persistent=False) -+ _validate_nonpersistent_grid_y(65_535, 1, is_persistent=False) -+ _validate_nonpersistent_grid_y(8_192, 8, is_persistent=True) -+ -+ with pytest.raises(ValueError, match=r"grid\.y would be 65536"): -+ _validate_nonpersistent_grid_y(8_192, 8, is_persistent=False) -+ -+ -+def test_mla_reducer_d_tile_selection(): -+ """Use output bands only when they shorten an underfilled reducer wave.""" -+ if not is_cute_dsl_available(): -+ pytest.skip("CuTe DSL not available") -+ -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ _get_reducer_d_tiles, -+ ) -+ -+ # B1/H32 and B1/H96 use D4; B1/H64 uses D2. H128 already fills a wave. -+ assert _get_reducer_d_tiles(1, 1, 32, 148, 32) == 4 -+ assert _get_reducer_d_tiles(1, 1, 64, 148, 32) == 2 -+ assert _get_reducer_d_tiles(1, 1, 96, 148, 32) == 4 -+ # Prefer the smaller tied topology and avoid duplication once rows fill a wave. -+ assert _get_reducer_d_tiles(1, 1, 48, 148, 32) == 2 -+ assert _get_reducer_d_tiles(1, 1, 128, 148, 32) == 1 -+ assert _get_reducer_d_tiles(4, 1, 96, 148, 32) == 1 -+ assert _get_reducer_d_tiles(1, 1, 96, 0, 32) == 1 -+ # Do not duplicate LSE work when a short sequence exposes too few splits. -+ assert _get_reducer_d_tiles(1, 1, 96, 148, 1) == 1 -+ assert _get_reducer_d_tiles(1, 1, 24, 148, 2) == 2 -+ -+ -+def test_mla_reducer_direct_class_capacity_defaults(): -+ """Direct class users keep the generic capacity unless opting into a cap.""" -+ if not is_cute_dsl_available(): -+ pytest.skip("CuTe DSL not available") -+ -+ import cutlass -+ -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode_fp16 import ( -+ BlackwellMultiHeadLatentAttentionForwardFP16, -+ ) -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode_fp8 import ( -+ BlackwellMultiHeadLatentAttentionForwardFP8, -+ ) -+ from flashinfer.cute_dsl.attention.monolithic.mla_helpers import MAX_SPLITS -+ -+ kwargs = dict( -+ acc_dtype=cutlass.Float32, -+ lse_dtype=cutlass.Float32, -+ mma_qk_tiler_mn=(128, 128), -+ mma_pv_tiler_mn=(128, 256), -+ max_active_clusters=1, -+ page_size=64, -+ skip_correction_threshold=0.0, -+ is_persistent=True, -+ is_var_seq=False, -+ is_var_split_kv=False, -+ enable_pdl=False, -+ ) -+ for kernel_cls in ( -+ BlackwellMultiHeadLatentAttentionForwardFP16, -+ BlackwellMultiHeadLatentAttentionForwardFP8, -+ ): -+ assert kernel_cls(**kwargs).reducer_max_splits == MAX_SPLITS -+ assert kernel_cls(**kwargs, reducer_max_splits=32).reducer_max_splits == 32 -+ with pytest.raises(ValueError, match="variable split-KV"): -+ kernel_cls(**{**kwargs, "is_var_split_kv": True}, reducer_max_splits=32) -+ -+ -+def test_flat_q_tile_split_workspace_geometry(): -+ """Workspace sizing uses the flattened M128 query-tile count.""" -+ if not is_cute_dsl_available(): -+ pytest.skip("CuTe DSL not available") -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ _get_split_kv_and_workspace_size, -+ ) -+ -+ batch_size, num_heads, q_len, num_q_tiles = 1, 48, 8, 3 -+ latent_dim = 512 -+ max_active_blocks = 148 -+ split_kv, workspace_size = _get_split_kv_and_workspace_size( -+ batch_size, -+ q_len, -+ num_heads, -+ latent_dim, -+ max_active_blocks, -+ ) -+ expected_split = min(max_active_blocks // (batch_size * num_q_tiles * 2), 32) -+ expected_workspace = ( -+ batch_size * 128 * num_q_tiles * expected_split * (latent_dim + 1) * 4 -+ ) -+ assert split_kv == expected_split -+ assert workspace_size == expected_workspace -+ -+ # Once one cluster per (batch, query tile) fills the machine, split-KV is -+ # disabled and no workspace is required. -+ direct_batch = max_active_blocks // (num_q_tiles * 2) -+ split_kv, workspace_size = _get_split_kv_and_workspace_size( -+ direct_batch, -+ q_len, -+ num_heads, -+ latent_dim, -+ max_active_blocks, -+ ) -+ assert split_kv == 1 -+ assert workspace_size == 0 -+ -+ -+def test_h96_sq8_split_workspace_drops_empty_partition(): -+ """Size only the eight nonempty K partitions at B1/H96/Sq8.""" -+ if not is_cute_dsl_available(): -+ pytest.skip("CuTe DSL not available") -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ _get_split_kv_and_workspace_size, -+ ) -+ -+ split_kv, workspace_size = _get_split_kv_and_workspace_size( -+ 1, 8, 96, 512, 148, 1024 -+ ) -+ assert split_kv == 8 -+ assert workspace_size == 1 * 128 * 6 * 8 * (512 + 1) * 4 -+ -+ -+def test_cute_dsl_workspace_sizer_follows_selected_impl(): -+ """Autotuning must size workspace with the implementation it will launch.""" -+ if not is_cute_dsl_available(): -+ pytest.skip("CuTe DSL not available") -+ -+ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( -+ _get_split_kv_and_workspace_size as monolithic_sizer, -+ ) -+ from flashinfer.cute_dsl.attention.wrappers.batch_mla import ( -+ _get_split_kv_and_workspace_size as modular_sizer, -+ ) -+ from flashinfer.mla._core import _get_cute_dsl_workspace_sizer -+ -+ assert _get_cute_dsl_workspace_sizer("monolithic", None) is monolithic_sizer -+ assert _get_cute_dsl_workspace_sizer("modular", None) is modular_sizer -+ assert _get_cute_dsl_workspace_sizer("auto", None) is monolithic_sizer -+ assert _get_cute_dsl_workspace_sizer("auto", torch.empty(0)) is modular_sizer -+ -+ -+def test_monolithic_workspace_cap_drops_empty_partitions(): -+ """Autotuning must use the launcher's max-sequence split normalization.""" -+ if not is_cute_dsl_available(): -+ pytest.skip("CuTe DSL not available") -+ -+ from flashinfer.mla._core import _cute_dsl_max_supported_batch -+ -+ # One K tile needs no split workspace for any candidate batch. Without -+ # max_seq_len propagation this is conservatively sized as 32 splits and -+ # the zero-byte workspace incorrectly caps the autotune sweep at B=1. -+ assert ( -+ _cute_dsl_max_supported_batch( -+ workspace_bytes=0, -+ q_len=1, -+ num_heads=128, -+ kv_lora_rank=512, -+ max_active_blocks=148, -+ max_seq_len=128, -+ candidate_max=8, -+ cute_dsl_impl="monolithic", -+ sinks=None, - ) -+ == 8 -+ ) - - - @pytest.mark.parametrize("batch_size", [1, 4, 16]) -diff --git a/tests/autotuner/test_autotuner_core.py b/tests/autotuner/test_autotuner_core.py -index a36c1e88..a5cf0100 100644 ---- a/tests/autotuner/test_autotuner_core.py -+++ b/tests/autotuner/test_autotuner_core.py -@@ -14,6 +14,7 @@ from flashinfer.fused_moe.utils import ( - make_hybrid_bucket_mapper, - ) - from flashinfer.mla._core import ( -+ CuteDslMlaDecodeRunner, - _build_mla_decode_tuning_config, - _mla_decode_tuning_config, - ) -@@ -1446,3 +1447,34 @@ def test_find_nearest_profile_cache_dedups_mla_decode_config(): - finally: - AutoTuner._find_nearest_profile.cache_clear() - _mla_decode_tuning_config.cache_clear() -+ -+ -+def _cute_dsl_runner_cache_extras(max_seq_len: int, workspace_bytes: int): -+ runner = object.__new__(CuteDslMlaDecodeRunner) -+ runner.kv_cache = torch.empty((1, 32, 576), dtype=torch.bfloat16) -+ runner.workspace_buffer = torch.empty(workspace_bytes, dtype=torch.uint8) -+ runner.qk_nope_head_dim = 512 -+ runner.kv_lora_rank = 512 -+ runner.qk_rope_head_dim = 64 -+ runner.page_size = 32 -+ runner.max_seq_len = max_seq_len -+ runner.is_var_seq = True -+ runner.uses_shared_paged_kv_idx = True -+ runner.enable_pdl = False -+ runner.sinks = None -+ runner.cute_dsl_impl = "auto" -+ runner._resolved_cute_dsl_impl = "monolithic" -+ -+ query = torch.empty((1, 1, 128, 576), dtype=torch.bfloat16) -+ out = torch.empty((1, 1, 128, 512), dtype=torch.bfloat16) -+ return runner.get_cache_key_extras([query, None, None, out]) -+ -+ -+def test_cute_dsl_runner_cache_tracks_split_workspace_geometry(): -+ """Cache hits must not bypass sequence- or capacity-dependent validity.""" -+ key_257 = _cute_dsl_runner_cache_extras(257, 1_000_000) -+ key_385 = _cute_dsl_runner_cache_extras(385, 1_000_000) -+ assert key_257 != key_385 -+ -+ key_small_workspace = _cute_dsl_runner_cache_extras(257, 800_000) -+ assert key_257 != key_small_workspace diff --git a/docker/kimi_k3/kimi_k3_cu12.Dockerfile b/docker/kimi_k3/kimi_k3_cu12.Dockerfile index 5b57ad607..b9eff4f37 100644 --- a/docker/kimi_k3/kimi_k3_cu12.Dockerfile +++ b/docker/kimi_k3/kimi_k3_cu12.Dockerfile @@ -4,15 +4,13 @@ # (deepseek-ai@d28bd67 at /sgl-workspace/DeepEP), the deep_gemm pip package, # and the CUDA 12.9 toolchain. # -# This image adds the four Kimi-K3-specific pieces that stock lacks: +# This image adds the three Kimi-K3-specific pieces that stock lacks: # 1. the Kimi-K3 SGLang code (this repo), editable-installed # 2. DeepEP patch + rebuild: # topk 11->16, SWITCH_HIDDEN += 3584, EP>8 SourceMeta alignment, # and cross-node timeout headroom; rebuilt for sm_90 and sm_100a only # 3. DeepGEMM upgrade to 0.1.5.post2: # official MegaMoE runtime-JIT header with Kimi-K3 SiTU support -# 4. FlashInfer CuTeDSL MLA DCP patch: -# apply the seven runtime-file diffs; exclude tests absent from the wheel # # Build (on/for x86_64; nvcc cross-compiles the DeepEP cubin, no GPU needed): # docker build -f docker/kimi_k3/kimi_k3_cu12.Dockerfile \ @@ -36,9 +34,7 @@ ENV RUSTUP_HOME="/usr/local/rustup" \ RUN apt-get update && \ apt-get install -y --no-install-recommends \ ca-certificates \ - curl \ - unzip \ - wget && \ + curl && \ curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ sh -s -- -y --no-modify-path --profile minimal \ --default-toolchain "${RUST_VERSION}" && \ @@ -62,12 +58,10 @@ RUN set -eu; \ # --- 1. Kimi-K3 SGLang code (replaces the base's stock sglang, editable) --- # Keep the installed extension modules, but discard Rust and pip build # artifacts that are not used at runtime. -ARG SGLANG_COMMIT="25035bff8d34f3fcce2c1a2a5b1fe610225e84ed" RUN rm -rf /sgl-workspace/sglang && \ - git clone --no-checkout \ + git clone --branch main \ https://github.com/sgl-project/sglang.git /sgl-workspace/sglang && \ cd /sgl-workspace/sglang && \ - git checkout --detach "${SGLANG_COMMIT}" && \ rm -rf .git && \ test ! -e .git && \ pip install -e python --no-deps && \ @@ -96,55 +90,24 @@ RUN python3 -m pip install \ "nvidia-nvimgcodec-cu12[all]==${NVIMGCODEC_VERSION}" && \ rm -rf /root/.cache/pip -# Install the pinned FlashInfer MXFP4 MoE runner cubin pool. -ARG TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip" -ARG TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192" -ARG TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1" -ENV SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool" - -RUN cubin_archive="/tmp/trtllm_gen_moe_cubin_pool.zip" && \ - cubin_extract_dir="/tmp/trtllm_gen_moe_cubin_extract" && \ - wget --no-verbose --output-document="${cubin_archive}" \ - "${TRTLLM_GEN_MOE_CUBIN_URL}" && \ - echo "${TRTLLM_GEN_MOE_CUBIN_SHA256} ${cubin_archive}" | \ - sha256sum --check --strict - && \ - mkdir -p "${cubin_extract_dir}" && \ - unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" && \ - test ! -e "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \ - mv "${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" \ - "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \ - test "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" \ - -type f -name '*.cubin' | wc -l)" -eq 1696 && \ - rm -f "${cubin_archive}" && \ - rm -rf "${cubin_extract_dir}" - -# Reinstall the matching FlashInfer package trio before patching its Python -# sources. A mixed Python/cubin/JIT-cache installation fails at import time. +# Install the matching official FlashInfer package trio. A mixed +# Python/cubin/JIT-cache installation fails at import time. # flashinfer-python and flashinfer-cubin are CUDA-independent packages; the # JIT-cache wheel is selected from the official CUDA 12.9 index. RUN python3 -m pip uninstall -y \ flashinfer-python flashinfer-cubin flashinfer-jit-cache && \ rm -rf /root/.cache/flashinfer /root/.cache/pip && \ python3 -m pip install --no-deps \ - "flashinfer-python==0.6.15.post1" && \ + "flashinfer-python==0.6.17" && \ python3 -m pip install --no-deps \ - "flashinfer-cubin==0.6.15.post1" \ + "flashinfer-cubin==0.6.17" \ --index-url https://flashinfer.ai/whl && \ python3 -m pip install --no-deps \ - "flashinfer-jit-cache==0.6.15.post1" \ + "flashinfer-jit-cache==0.6.17" \ --index-url https://flashinfer.ai/whl/cu129 && \ - python3 -c 'from importlib.metadata import version; expected = "0.6.15.post1"; assert version("flashinfer-python").split("+", 1)[0] == expected; assert version("flashinfer-cubin").split("+", 1)[0] == expected; assert version("flashinfer-jit-cache").startswith(expected + "+cu129"), version("flashinfer-jit-cache")' && \ + python3 -c 'from importlib.metadata import version; expected = "0.6.17"; assert version("flashinfer-python").split("+", 1)[0] == expected; assert version("flashinfer-cubin").split("+", 1)[0] == expected; assert version("flashinfer-jit-cache").startswith(expected + "+cu129"), version("flashinfer-jit-cache")' && \ rm -rf /root/.cache/pip -ENV FLASHINFER_VERSION="0.6.15.post1" - -# --- 4. FlashInfer: CuTeDSL MLA decode-context-parallel runtime patch --- -RUN FLASHINFER_DCP_PATCH=/sgl-workspace/sglang/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt && \ - FLASHINFER_SITE_PACKAGES="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')" && \ - sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \ - patch --dry-run --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \ - sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \ - patch --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \ - rm -rf /root/.cache/flashinfer /root/.cache/pip +ENV FLASHINFER_VERSION="0.6.17" WORKDIR /sgl-workspace/sglang diff --git a/docker/kimi_k3/kimi_k3_cu13.Dockerfile b/docker/kimi_k3/kimi_k3_cu13.Dockerfile index 58b884ce4..ed738bab4 100644 --- a/docker/kimi_k3/kimi_k3_cu13.Dockerfile +++ b/docker/kimi_k3/kimi_k3_cu13.Dockerfile @@ -4,7 +4,7 @@ # (deepseek-ai@d28bd67 at /sgl-workspace/DeepEP), the deep_gemm pip package, # and the CUDA 13 toolchain (nvcc + /usr/local/cuda/include/cccl). # -# This image adds the four Kimi-K3-specific pieces that stock lacks: +# This image adds the three Kimi-K3-specific pieces that stock lacks: # 1. the Kimi-K3 SGLang code (this repo), editable-installed # 2. DeepEP patch + rebuild: # topk 11->16, SWITCH_HIDDEN += 3584, EP>8 SourceMeta alignment, @@ -12,8 +12,6 @@ # sm_90, sm_100a, and sm_103a # 3. DeepGEMM upgrade to 0.1.5.post2: # official MegaMoE runtime-JIT header with Kimi-K3 SiTU support -# 4. FlashInfer CuTeDSL MLA DCP patch: -# apply the seven runtime-file diffs; exclude tests absent from the wheel # # Build (on/for aarch64; nvcc cross-compiles the DeepEP cubin, no GPU needed): # docker build -f docker/kimi_k3/kimi_k3_cu13.Dockerfile \ @@ -37,9 +35,7 @@ ENV RUSTUP_HOME="/usr/local/rustup" \ RUN apt-get update && \ apt-get install -y --no-install-recommends \ ca-certificates \ - curl \ - unzip \ - wget && \ + curl && \ curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ sh -s -- -y --no-modify-path --profile minimal \ --default-toolchain "${RUST_VERSION}" && \ @@ -53,12 +49,10 @@ ARG TORCH_CUDA_ARCH_LIST="9.0;10.0a;10.3a" # --- 1. Kimi-K3 SGLang code (replaces the base's stock sglang, editable) --- # Keep the installed extension modules, but discard Rust and pip build # artifacts that are not used at runtime. -ARG SGLANG_COMMIT="25035bff8d34f3fcce2c1a2a5b1fe610225e84ed" RUN rm -rf /sgl-workspace/sglang && \ - git clone --no-checkout \ + git clone --branch main \ https://github.com/sgl-project/sglang.git /sgl-workspace/sglang && \ cd /sgl-workspace/sglang && \ - git checkout --detach "${SGLANG_COMMIT}" && \ rm -rf .git && \ test ! -e .git && \ pip install -e python --no-deps && \ @@ -86,53 +80,22 @@ RUN python3 -m pip install \ "nvidia-nvimgcodec-cu13[all]==${NVIMGCODEC_VERSION}" && \ rm -rf /root/.cache/pip -# Install the pinned FlashInfer MXFP4 MoE runner cubin pool. -ARG TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip" -ARG TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192" -ARG TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1" -ENV SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool" - -RUN cubin_archive="/tmp/trtllm_gen_moe_cubin_pool.zip" && \ - cubin_extract_dir="/tmp/trtllm_gen_moe_cubin_extract" && \ - wget --no-verbose --output-document="${cubin_archive}" \ - "${TRTLLM_GEN_MOE_CUBIN_URL}" && \ - echo "${TRTLLM_GEN_MOE_CUBIN_SHA256} ${cubin_archive}" | \ - sha256sum --check --strict - && \ - mkdir -p "${cubin_extract_dir}" && \ - unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" && \ - test ! -e "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \ - mv "${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" \ - "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \ - test "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" \ - -type f -name '*.cubin' | wc -l)" -eq 1696 && \ - rm -f "${cubin_archive}" && \ - rm -rf "${cubin_extract_dir}" - -# Reinstall the matching FlashInfer package trio before patching its Python -# sources. A mixed Python/cubin/JIT-cache installation fails at import time. +# Install the matching official FlashInfer package trio. A mixed +# Python/cubin/JIT-cache installation fails at import time. RUN python3 -m pip uninstall -y \ flashinfer-python flashinfer-cubin flashinfer-jit-cache && \ rm -rf /root/.cache/flashinfer /root/.cache/pip && \ python3 -m pip install --no-deps \ - "flashinfer-python==0.6.15.post1" && \ + "flashinfer-python==0.6.17" && \ python3 -m pip install --no-deps \ - "flashinfer-cubin==0.6.15.post1" \ + "flashinfer-cubin==0.6.17" \ --index-url https://flashinfer.ai/whl && \ python3 -m pip install --no-deps \ - "flashinfer-jit-cache==0.6.15.post1" \ + "flashinfer-jit-cache==0.6.17" \ --index-url https://flashinfer.ai/whl/cu130 && \ - python3 -c 'from importlib.metadata import version; expected = "0.6.15.post1"; packages = ("flashinfer-python", "flashinfer-cubin", "flashinfer-jit-cache"); actual = {package: version(package).split("+", 1)[0] for package in packages}; assert all(value == expected for value in actual.values()), actual' && \ + python3 -c 'from importlib.metadata import version; expected = "0.6.17"; packages = ("flashinfer-python", "flashinfer-cubin", "flashinfer-jit-cache"); actual = {package: version(package).split("+", 1)[0] for package in packages}; assert all(value == expected for value in actual.values()), actual' && \ rm -rf /root/.cache/pip -ENV FLASHINFER_VERSION="0.6.15.post1" - -# --- 4. FlashInfer: CuTeDSL MLA decode-context-parallel runtime patch --- -RUN FLASHINFER_DCP_PATCH=/sgl-workspace/sglang/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt && \ - FLASHINFER_SITE_PACKAGES="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')" && \ - sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \ - patch --dry-run --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \ - sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \ - patch --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \ - rm -rf /root/.cache/flashinfer /root/.cache/pip +ENV FLASHINFER_VERSION="0.6.17" WORKDIR /sgl-workspace/sglang diff --git a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx index a3be9cf21..b255fe19a 100644 --- a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx +++ b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx @@ -127,16 +127,7 @@ Capacity levers, all in the Playground. Each trades precision or cache behavior Speculation: DSPARK holds block size + 1 (= 8) intermediate states per request — the calculator folds this in — and an unset `--max-running-requests` resets to 48 under spec (the command panel reminds you; set it explicitly to raise). -**MoE runner.** Leave `--moe-runner-backend` unset on Blackwell and it resolves to FlashInfer MXFP4 (W4A8, prebuilt trtllm-gen SiTU kernels) when the cubin pool is installed, Marlin (W4A16) otherwise; H100/H200 pin Marlin. The B200 Balanced and High-Throughput cells pin `flashinfer_mxfp4` explicitly because that is the shape they were brought up on — on an install without the pool, drop the flag to fall back to Marlin. The published Docker images already provision the **SiTU cubin pool**; to install it independently, run the same flow as the Dockerfile: - -```bash -wget https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip -sudo mkdir -p /opt/trtllm_gen_moe_cubin_pool -sudo unzip -q trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip -d /opt/trtllm_gen_moe_cubin_pool -export SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=/opt/trtllm_gen_moe_cubin_pool/trtllm_gen_moe_cubin_pool_20260617_v0613rc1 -``` - -Remaining kernel sources JIT once from the public `flashinfer` wheel (a few minutes, cached). +**MoE runner.** Leave `--moe-runner-backend` unset on Blackwell: FlashInfer MXFP4 (W4A8, official trtllm-gen SiTU kernels) is selected with the pinned FlashInfer 0.6.17 dependency; H100/H200 pin Marlin. The B200 Balanced and High-Throughput cells pin `flashinfer_mxfp4` explicitly because that is the shape they were brought up on. The published Docker images install the matching official `flashinfer-python`, `flashinfer-cubin`, and `flashinfer-jit-cache` packages. **Attention backend.** Leave all three attention knobs unset on Blackwell: K3 resolves prefill, decode, and — under DSPARK — verification as a set (`trtllm_mla` across the board; `cutedsl_mla` takes decode and verification under DCP). On the non-DCP recipes, setting any one of the three cancels the auto-resolution for the others. The B200 Balanced and High-Throughput cells pin `--decode-attention-backend cutedsl_mla`, which is what auto-resolution picks for those DCP recipes anyway — it is written out because it is the shape they were brought up on, not because it changes the resolution. H100/H200 pin `flashmla` for decode. @@ -395,7 +386,7 @@ Both presets are one click away in the [Playground above](#playground): pick a * Decisions the preset already makes: -- **MegaMoE on `deep_gemm`** — the fastest a2a backend; needs the SiTU cubin pool ([§2](#2-configuration-tips)). +- **MegaMoE + `deep_gemm`** — the fused DeepGEMM all-to-all/MoE path used by these large-scale DP/EP throughput presets, with K3's SiTU activation. - **SP-MoE and shared-expert overlap** engage automatically under EP a2a; the K3 all-reduce fusion does not. - **Spec Decode follows the Deploy knob.** Acceptance thins at large batch; spec × EP × DP-attention is validated only at 8-GPU EP8 × DP2 (full GSM8K) — experimental at these scales. diff --git a/docs/src/snippets/_playground.jsx b/docs/src/snippets/_playground.jsx index 45f1a5c41..a7ccb4e7e 100644 --- a/docs/src/snippets/_playground.jsx +++ b/docs/src/snippets/_playground.jsx @@ -629,7 +629,7 @@ export const Playground = ({ config }) => { "--moe-a2a-backend", "--moe-runner-backend", ]); // Backend options may carry their own env (e.g. the FlashInfer MXFP4 - // cubin-pool path): strip every backend option's env keys, then + // backend-specific path): strip every backend option's env keys, then // re-add the selected option's. const backendEnvKeys = []; for (const o of (fc.backend?.options || [])) { diff --git a/docs/src/snippets/configs/moonshotai/kimi-k3.jsx b/docs/src/snippets/configs/moonshotai/kimi-k3.jsx index 25ab622a5..309c86988 100644 --- a/docs/src/snippets/configs/moonshotai/kimi-k3.jsx +++ b/docs/src/snippets/configs/moonshotai/kimi-k3.jsx @@ -458,10 +458,8 @@ export const config = { // Blackwell-only kernel-fusion path; selecting it reveals the Quantization sub-select. { id: "megamoe", label: "MegaMoE", flags: ["--moe-a2a-backend megamoe"], requiresHw: ["b200", "b300", "gb200", "gb300"] }, - // Blackwell-only: runs the prebuilt trtllm-gen SiTU cubins; needs the - // downloadable SiTU cubin pool unpacked and pointed to by the env var. + // Blackwell-only: runs FlashInfer's official trtllm-gen SiTU kernels. { id: "flashinfer_mxfp4", label: "FlashInfer (MXFP4)", flags: ["--moe-runner-backend flashinfer_mxfp4"], - env: ["SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=/path/to/trtllm_gen_moe_cubin_pool"], requiresHw: ["b200", "b300", "gb200", "gb300"] }, { id: "marlin", label: "Marlin (W4A16)", flags: ["--moe-runner-backend marlin"] }, ], @@ -931,10 +929,8 @@ export const config = { "--pp-size 2", "--dcp-size 8", "--ep-size 8", - // Both pinned to the brought-up shape rather than left to the auto - // resolution the rest of Blackwell uses. The MXFP4 runner needs the SiTU - // cubin pool the published image ships; drop it to get the Marlin - // fallback on an install without one. + // Both backends are pinned to the validated recipe even though automatic + // resolution selects them on Blackwell. "--moe-runner-backend flashinfer_mxfp4", "--decode-attention-backend cutedsl_mla", "--mem-fraction-static 0.85", diff --git a/python/pyproject.toml b/python/pyproject.toml index d093c365b..487213fca 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "einops", "fastapi", "flash-attn-4>=4.0.0b18", - "flashinfer_python[cu13]==0.6.15.post1", # keep it aligned with jit-cache version in Dockerfile + "flashinfer_python[cu13]==0.6.17", # keep it aligned with jit-cache version in Dockerfile "gguf", "humming-kernels[cu13]==0.1.10", "interegular", diff --git a/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh index 55ec8fafc..26ebd56e0 100644 --- a/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh +++ b/python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh @@ -152,7 +152,7 @@ __global__ __launch_bounds__(1024, 1) void all_reduce_push_res_kernel(const __gr // --- deferred-finalize staging (finalize_push_norm) ------------------------ // The trtllm-gen MoE with do_finalize=False hands back its finalize inputs -// (see kernels/ops/moe/trtllm_gen_moe.py); the fused kernel computes the finalize +// (see FlashInfer's TRT-LLM-gen MoE); the fused kernel computes the finalize // during the push staging pass, so the rank-local latent never materializes. constexpr uint32_t kFinTopK = 16; diff --git a/python/sglang/kernels/ops/moe/trtllm_gen_moe.py b/python/sglang/kernels/ops/moe/trtllm_gen_moe.py deleted file mode 100644 index c2da4402d..000000000 --- a/python/sglang/kernels/ops/moe/trtllm_gen_moe.py +++ /dev/null @@ -1,528 +0,0 @@ -"""TRT-LLM-gen fused MoE (SiTU) compiled through the sglang JIT system. - -Builds the trtllm-gen fused-MoE host/runner sources with sglang's own -tvm-ffi ``load_jit`` from a **self-contained cubin pool** -(``SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL``): a downloadable directory holding - - * the prebuilt SiTU cubins (``local/``) + ``config.json`` + - ``flashinferMetaInfo.h``, - * the flat batched-gemm ABI headers (staged into a - ``trtllmGen_bmm_export/``-shaped include tree at build time), - * an ``overlay/`` with only the sources/headers that differ from the - public ``flashinfer`` pip package. - -Every unmodified source and the CUTLASS headers come from the installed -``flashinfer`` package's ``data/`` tree (the wheel ships it for its own -JIT), so running this backend needs exactly one download and one env var — -no extra source checkout. - -This module vendors only glue: - - * header staging: the pool ships the batched-gemm ABI headers flat; they - are copied into a content-addressed include tree shaped like - ``flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/``; - * JIT build of the 12 launcher/runner/routing sources with the private - ABI defines (``TLLM_GEN_LOCAL_CUBINS_ABI`` etc.); - * the ctypes cubin-loader callback (the .so asks for cubins by absolute - path + sha256; we read them from the pool); - * a thin ``trtllm_fp4_block_scale_moe`` wrapper (FromLogits routing, - ``do_finalize=True``); kernel tile config ("tactic") defaults to the - runner's built-in heuristic — pass an explicit one for tuned setups. - -Validated for the Kimi K3 decode/prefill MoE regime: MxFP4 weights with -bf16 (w4a16) or MxFP8 (w4a8) activations, ``ActivationType.Situ`` (SiTuGlu: -``a*tanh(g/a)*sigmoid(g) * b*tanh(u/b)``), DeepSeekV3/noaux_tc routing. -""" - -from __future__ import annotations - -import ctypes -import hashlib -import logging -import os -import pathlib -import shutil -from typing import TYPE_CHECKING, Optional, Sequence - -import torch - -from sglang.kernels.jit.utils import ( - cache_once, - get_jit_cuda_arch, - load_jit, - override_jit_cuda_arch, -) -from sglang.srt.environ import envs - -if TYPE_CHECKING: - from tvm_ffi.module import Module - -# ActivationType / RoutingMethodType values from trtllm-gen's tllm_enums -# (kept as plain ints here to avoid importing anything for them). -ACTIVATION_SITU = 9 -ROUTING_DEEPSEEK_V3 = 2 -_ROUTING_TOPK = 5 -_ROUTING_INPUT_FROM_LOGITS = 0 -# NOTE: the enum VALUES start at 0; the "Mode 1/2/3" wording in upstream -# comments is documentation numbering, not the enum value. -_ROUTING_INPUT_PACKED = 1 - -# Batched-gemm ABI headers shipped flat in the cubin pool; the launcher -# includes them as flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/. -_BMM_EXPORT_HEADERS = [ - "BatchedGemmEnums.h", - "BatchedGemmInterface.h", - "BatchedGemmOptions.h", - "Enums.h", - "GemmGatedActOptions.h", - "GemmOptions.h", - "KernelParams.h", - "KernelParamsDecl.h", - "KernelTraits.h", - "TmaDescriptor.h", - "trtllm/gen/CommonUtils.h", - "trtllm/gen/CudaArchDecl.h", - "trtllm/gen/CudaKernelLauncher.h", - "trtllm/gen/DtypeDecl.h", - "trtllm/gen/MmaDecl.h", - "trtllm/gen/SfLayoutDecl.h", - "trtllm/gen/SparsityDecl.h", -] - -_SOURCES = [ - "csrc/nv_internal/cpp/kernels/quantization.cu", - "csrc/nv_internal/cpp/common/envUtils.cpp", - "csrc/nv_internal/cpp/common/logger.cpp", - "csrc/nv_internal/cpp/common/stringUtils.cpp", - "csrc/nv_internal/cpp/common/tllmException.cpp", - "csrc/nv_internal/cpp/common/memoryUtils.cu", - "csrc/trtllm_fused_moe_kernel_launcher.cu", - "csrc/trtllm_fused_moe_runner.cu", - "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_deepseek.cu", - "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_llama4.cu", - "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu", - "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_common.cu", - "csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu", - "csrc/trtllm_batched_gemm_runner.cu", -] - - -logger = logging.getLogger(__name__) - - -def cubin_pool_dir() -> Optional[pathlib.Path]: - p = envs.SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL.get() - if not p: - return None - pool = pathlib.Path(p) - return pool if pool.is_dir() else None - - -def _flashinfer_data_dir() -> Optional[pathlib.Path]: - """The installed public flashinfer package's JIT source tree (ships - csrc/, include/ and its pinned cutlass), used as the base layer under - the pool's overlay.""" - try: - import flashinfer # noqa: PLC0415 - except ImportError: - return None - data = pathlib.Path(flashinfer.__file__).parent / "data" - return data if (data / "csrc").is_dir() else None - - -def available() -> bool: - pool = cubin_pool_dir() - return ( - pool is not None - and (pool / "flashinferMetaInfo.h").is_file() - and (pool / "local").is_dir() - # Modified sources ship in the pool's overlay/, everything else - # compiles from the installed flashinfer package. - and (pool / "overlay" / "csrc").is_dir() - and _flashinfer_data_dir() is not None - ) - - -def _stage_headers(pool: pathlib.Path) -> pathlib.Path: - """Copy the pool's ABI headers into a content-addressed include tree.""" - meta = (pool / "flashinferMetaInfo.h").read_bytes() - tag = hashlib.sha256(meta).hexdigest()[:12] - cache = pathlib.Path( - os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi") - ).expanduser() - root = cache / "trtllm_gen_moe_headers" / tag - dest = root / "flashinfer" / "trtllm" / "batched_gemm" / "trtllmGen_bmm_export" - stamp = root / ".staged" - if not stamp.is_file(): - for name in _BMM_EXPORT_HEADERS: - target = dest / name - target.parent.mkdir(parents=True, exist_ok=True) - shutil.copyfile(pool / name, target) - shutil.copyfile(pool / "flashinferMetaInfo.h", dest / "flashinferMetaInfo.h") - stamp.touch() - return root - - -def _cuda_home() -> pathlib.Path: - home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") - if not home: - nvcc = shutil.which("nvcc") - home = str(pathlib.Path(nvcc).parent.parent) if nvcc else "/usr/local/cuda" - return pathlib.Path(home) - - -def _cuda_include_dir() -> str: - return str(_cuda_home() / "include") - - -def _cuda_stub_ldflags() -> list[str]: - """-L flags for the libcuda driver stub, so -lcuda links in bare build - environments (containers without the driver lib on the default linker - path); the real driver is dlopened at runtime as usual.""" - home = _cuda_home() - stubs = [ - home / "lib64" / "stubs", - *home.glob("targets/*/lib/stubs"), - ] - return [f"-L{s}" for s in stubs if s.is_dir()] - - -_CUBIN_CB_KEEPALIVE = {} - - -def _setup_cubin_loader(so_path: str, pool_local: pathlib.Path) -> None: - """Register the ctypes callback the .so uses to fetch cubins by name. - - The runner requests ``/`` (absolute, - because the pool path is baked in at compile time); we read the bytes - and hand them back via FlashInferSetCurrentCubin. - """ - if so_path in _CUBIN_CB_KEEPALIVE: - return - lib = ctypes.CDLL(so_path) - cb_type = ctypes.CFUNCTYPE(None, ctypes.c_char_p, ctypes.c_char_p) - - def _get_cubin(name: bytes, sha256: bytes) -> None: - rel = name.decode() - path = pathlib.Path(rel) - if not path.is_absolute(): - path = pool_local / rel - if path.suffix != ".cubin": - path = path.with_name(path.name + ".cubin") - data = path.read_bytes() - want = sha256.decode() - if want: - got = hashlib.sha256(data).hexdigest() - if got != want: - raise RuntimeError( - f"cubin sha mismatch for {path}: want {want} got {got}" - ) - lib.FlashInferSetCurrentCubin( - ctypes.cast(ctypes.create_string_buffer(data, len(data)), ctypes.c_char_p), - ctypes.c_int(len(data)), - ) - - cb = cb_type(_get_cubin) - _CUBIN_CB_KEEPALIVE[so_path] = (lib, cb) - lib.FlashInferSetCubinCallback(cb) - - -@cache_once -def _jit_trtllm_gen_moe_module() -> Module: - pool = cubin_pool_dir() - fi_data = _flashinfer_data_dir() - if pool is None or not (pool / "overlay" / "csrc").is_dir() or fi_data is None: - raise RuntimeError( - "trtllm-gen MoE sources not found: point " - "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL at an unpacked cubin pool " - "(cubins + flat ABI headers + overlay/) and install the public " - "flashinfer package." - ) - # Overlay first (its modified sources/headers shadow the public copies), - # installed flashinfer data as the base. - src_roots = [pool / "overlay", fi_data] - include_roots = [pool / "overlay", fi_data] - - def _resolve_source(rel: str) -> str: - for root in src_roots: - cand = root / rel - if cand.is_file(): - return str(cand) - raise RuntimeError(f"trtllm-gen MoE source not found in any root: {rel}") - - staged = _stage_headers(pool) - meta_tag = staged.name - cubin_path = str((pool / "local").resolve()) - - cache = pathlib.Path( - os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi") - ).expanduser() - # Flags are not part of load_jit's source hash: fold the pool identity - # (meta hash + path) into the module marker so a pool change rebuilds. - path_tag = hashlib.sha256(cubin_path.encode()).hexdigest()[:8] - build_dir = cache / f"sgl_trtllm_gen_moe_{meta_tag}_{path_tag}" - - cpp_files = [_resolve_source(s) for s in _SOURCES if s.endswith(".cpp")] - cuda_files = [_resolve_source(s) for s in _SOURCES if s.endswith(".cu")] - # quantization.cu emits fp4 cvt instructions (.e2m1x2) that need the - # arch-specific feature set: compile for sm_XXXa, not plain sm_XXX. - # The trtllm-gen cubins themselves are prebuilt (sm100f) and loaded at - # runtime, unaffected by this flag. - arch = get_jit_cuda_arch() - with override_jit_cuda_arch(arch.major, arch.minor, "a"): - module = load_jit( - "trtllm_gen_moe", - meta_tag, - path_tag, - external_cpp_files=cpp_files, - external_cuda_files=cuda_files, - header_only=False, # the launcher exports its own tvm-ffi functions - extra_cflags=["-fvisibility=hidden"], - extra_cuda_cflags=[ - "-DTLLM_GEN_EXPORT_INTERFACE", - "-DTLLM_GEN_EXPORT_FLASHINFER", - "-DTLLM_ENABLE_CUDA", - "-DENABLE_BF16", - "-DENABLE_FP8", - "-DENABLE_FP4", - "-DCUTLASS_ENABLE_GDC_FOR_SM100=1", - "-DTLLM_GEN_LOCAL_CUBINS_ABI", - "-DFLASHINFER_PRIVATE_MOE_FFI_NAMES", - "-DFLASHINFER_PRIVATE_MOE_LEAN_ROUTING", - f'-DTLLM_GEN_GEMM_CUBIN_PATH=\\"{cubin_path}\\"', - "-Xcompiler=-fvisibility=hidden", - ], - extra_ldflags=[*_cuda_stub_ldflags(), "-lcuda", "-lnvrtc"], - extra_include_paths=[ - str(staged), - str( - staged - / "flashinfer" - / "trtllm" - / "batched_gemm" - / "trtllmGen_bmm_export" - ), - # Per-root include layout: include/, csrc/, csrc/nv_internal/, - # csrc/nv_internal/include/, plus the flashinfer package's - # pinned CUTLASS (data/cutlass/). The overlay root comes first - # so modified headers shadow the public copies. - *[ - str(root / sub) - for root in include_roots - for sub in ( - "include", - "csrc", - "csrc/nv_internal", - "csrc/nv_internal/include", - ) - ], - *[ - str(root / "cutlass" / "include") - for root in include_roots - if (root / "cutlass" / "include").is_dir() - ], - # Host .cpp files (g++) need the CUDA headers explicitly; nvcc - # adds them implicitly for .cu. CUDA 13's bundled CCCL is - # used as-is (mixing another pinned CCCL with the toolkit's - # explodes). - _cuda_include_dir(), - ], - build_directory=str(build_dir), - ) - so_files = sorted(build_dir.glob("*.so")) - if not so_files: - raise RuntimeError(f"no built .so under {build_dir}") - _setup_cubin_loader(str(so_files[-1]), pool / "local") - return module - - -def trtllm_fp4_block_scale_moe( - routing_logits: torch.Tensor, - routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, - hidden_states_scale: Optional[torch.Tensor], - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm1_alpha: Optional[torch.Tensor], - gemm1_beta: Optional[torch.Tensor], - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - output1_scale_scalar: Optional[torch.Tensor], - output1_scale_gate_scalar: Optional[torch.Tensor], - output2_scale_scalar: Optional[torch.Tensor], - num_experts: int, - top_k: int, - n_group: Optional[int], - topk_group: Optional[int], - intermediate_size: int, - routed_scaling_factor: Optional[float], - routing_method_type: int = ROUTING_DEEPSEEK_V3, - activation_type: int = ACTIVATION_SITU, - norm_topk_prob: bool = True, - local_expert_offset: int = 0, - local_num_experts: Optional[int] = None, - tactic: Sequence[int] = (-1, -1), - output: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """FP4 block-scale MoE with routing from logits and finalize fused. - - ``hidden_states``: bf16 ``[T, hidden]`` (w4a16) or MxFP8-packed uint8 - with ``hidden_states_scale`` (w4a8). Weights are trtllm-gen shuffled - MxFP4 (uint8 packed, fp8 block scales, MajorK). ``tactic`` is the - (gemm1, gemm2) config index pair; ``(-1, -1)`` = runner heuristic. - """ - module = _jit_trtllm_gen_moe_module() - # The FFI launcher reads these as dense row-major; a strided slice - # (e.g. a fused-GEMM split) would silently mis-route. - routing_logits = routing_logits.contiguous() - hidden_states = hidden_states.contiguous() - num_tokens = routing_logits.shape[0] - hidden_size = hidden_states.shape[-1] - if hidden_states.dtype == torch.uint8: - hidden_size *= 2 - device = hidden_states.device - topk_ids = torch.empty(num_tokens, top_k, dtype=torch.int32, device=device) - topk_weights = torch.empty( - num_tokens, top_k, dtype=routing_logits.dtype, device=device - ) - if output is None: - output = torch.empty( - num_tokens, hidden_size, dtype=torch.bfloat16, device=device - ) - module.trtllm_fp4_block_scale_moe_private( - _ROUTING_INPUT_FROM_LOGITS, - routing_logits, - topk_ids, - topk_weights, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - None, # gemm1_bias - gemm1_alpha, - gemm1_beta, - None, # gemm1_clamp_limit - gemm2_weights, - gemm2_weights_scale, - None, # gemm2_bias - output1_scale_scalar, - output1_scale_gate_scalar, - output2_scale_scalar, - None, # per_token_scale - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - num_experts if local_num_experts is None else local_num_experts, - routed_scaling_factor, - routing_method_type, - True, # do_finalize - True, # enable_pdl - activation_type, - output, - list(tactic), - norm_topk_prob, - None, # routing_replay_out - ) - return output - - -def trtllm_fp4_block_scale_routed_moe( - packed_topk_ids: torch.Tensor, - hidden_states: torch.Tensor, - hidden_states_scale: Optional[torch.Tensor], - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm1_alpha: Optional[torch.Tensor], - gemm1_beta: Optional[torch.Tensor], - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - output1_scale_scalar: Optional[torch.Tensor], - output1_scale_gate_scalar: Optional[torch.Tensor], - output2_scale_scalar: Optional[torch.Tensor], - num_experts: int, - top_k: int, - intermediate_size: int, - activation_type: int = ACTIVATION_SITU, - local_expert_offset: int = 0, - local_num_experts: Optional[int] = None, - tactic: Sequence[int] = (-1, -1), - output: Optional[torch.Tensor] = None, - do_finalize: bool = True, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """FP4 block-scale MoE with PRECOMPUTED routing (PackedPrecomputed). - - ``packed_topk_ids``: int32 ``[T, top_k]`` with ``(expert_id << 16) | - bf16-weight-bits`` (PackTopkIds layout) — selection and weights come - from the caller's router, the in-op routing kernels are skipped. This - is the fast path at small T, where the in-op single-CTA routing kernel - (~22 µs at 896 experts) costs more than an external radix router. - - ``do_finalize=False`` skips the in-op finalize (top-k weighted - unpermute) and returns its inputs instead: - ``(gemm2_output [padded_rows, hidden] bf16 in permuted layout, - topk_weights [T, top_k] bf16 unpacked from packed_topk_ids, - expanded_idx_to_permuted_idx [T*top_k] int32 with -1 = dropped slot)``. - ``output`` is left unwritten in that mode. - """ - module = _jit_trtllm_gen_moe_module() - hidden_states = hidden_states.contiguous() - num_tokens = packed_topk_ids.shape[0] - hidden_size = hidden_states.shape[-1] - if hidden_states.dtype == torch.uint8: - hidden_size *= 2 - device = hidden_states.device - # Mode 2 unpacks the weights in-kernel; this is its output buffer. - topk_weights = torch.empty(num_tokens, top_k, dtype=torch.bfloat16, device=device) - if output is None: - output = torch.empty( - num_tokens, hidden_size, dtype=torch.bfloat16, device=device - ) - result = module.trtllm_fp4_block_scale_moe_private( - _ROUTING_INPUT_PACKED, - None, # routing_logits - packed_topk_ids.contiguous(), - topk_weights, - None, # routing_bias (already applied by the external router) - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - None, # gemm1_bias - gemm1_alpha, - gemm1_beta, - None, # gemm1_clamp_limit - gemm2_weights, - gemm2_weights_scale, - None, # gemm2_bias - output1_scale_scalar, - output1_scale_gate_scalar, - output2_scale_scalar, - None, # per_token_scale - num_experts, - top_k, - None, # n_group - None, # topk_group - intermediate_size, - local_expert_offset, - num_experts if local_num_experts is None else local_num_experts, - 1.0, # routed_scaling_factor (already applied by the router) - _ROUTING_TOPK, # routing_method_type (unused for precomputed) - do_finalize, - True, # enable_pdl - activation_type, - output, - list(tactic), - True, # norm_topk_prob (unused for precomputed) - None, # routing_replay_out - ) - if do_finalize: - return output - # Deferred: [gemm2_output, expert_weights (None in packed mode — the - # weights live in the topk_weights buffer mode 2 unpacked into), - # expanded_idx_to_permuted_idx]. Index access — iterating the tvm-ffi - # Array yields one-shot dlpack capsules. - return result[0], topk_weights, result[2] diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 73282da1e..e665e8a25 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -371,13 +371,6 @@ def _dspark_verify_on_decode_backend( return False -_KIMI_K3_DCP_PATCH_URL = ( - "https://github.com/sgl-project/sglang/blob/" - "b701464720ca22aa1851d5dda7144e84a410f2c7/" - "docker/kimi_k3/kimi_k3_cu13.Dockerfile#L116-L123" -) - - def _require_kimi_k3_cutedsl_dcp_support() -> None: try: from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla @@ -386,17 +379,16 @@ def _require_kimi_k3_cutedsl_dcp_support() -> None: except (ImportError, TypeError, ValueError) as exc: raise RuntimeError( "Kimi-K3 DCP with decode_attention_backend='cutedsl_mla' requires " - "a DCP-patched FlashInfer " - "trtllm_batch_decode_with_kv_cache_mla exposing enable_dcp in its " - f"signature. Apply the patch as shown in {_KIMI_K3_DCP_PATCH_URL}." + "FlashInfer 0.6.17 or newer with " + "trtllm_batch_decode_with_kv_cache_mla exposing enable_dcp." ) from exc if "enable_dcp" not in parameters: raise RuntimeError( "Kimi-K3 DCP with decode_attention_backend='cutedsl_mla' requires " "enable_dcp in the signature of " - "flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla. Apply " - f"the FlashInfer DCP patch as shown in {_KIMI_K3_DCP_PATCH_URL}." + "flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla; upgrade " + "to FlashInfer 0.6.17 or newer." ) @@ -546,43 +538,19 @@ def _kimi_k3_moe_runner_overrides(server_args: Any, hf_config: Any) -> dict: # MoE runner default, independent of the attention-backend gate above. # trtllm-gen fused MoE (flashinfer_mxfp4) beats marlin on both the decode # (M=bs) and the target-verify (M=bs*(gamma+1)) regimes on SM100/SM103; - # it hard-requires the SiTU cubin pool on the box (K3's SiTU activation has - # no public cubins). Do not silently trade W4A8 for Marlin W4A16 when the - # default cannot start; explicit non-FlashInfer runner choices still win. - if server_args.moe_runner_backend not in ("auto", "flashinfer_mxfp4"): + # FlashInfer 0.6.17+ ships the required SiTU kernels and is a pinned + # project dependency. + if server_args.moe_runner_backend != "auto": return {} if not (is_sm100_supported() and get_device_sm() in (100, 103)): return {} if not _is_mxfp4_pack_quantized(hf_config): return {} - from sglang.kernels.ops.moe.trtllm_gen_moe import available as _trtllm_gen_moe_ok - - if not _trtllm_gen_moe_ok(): - raise RuntimeError( - "Kimi-K3 on Blackwell with moe_runner_backend='auto' or " - "'flashinfer_mxfp4' requires a valid " - "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL. Install it with:\n" - "wget https://github.com/sgl-project/whl/releases/download/" - "trtllm_gen_moe_cubin_20260617/" - "trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip\n" - "sudo mkdir -p /opt/trtllm_gen_moe_cubin_pool\n" - "sudo unzip -q " - "trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip -d " - "/opt/trtllm_gen_moe_cubin_pool\n" - "export " - "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=/opt/trtllm_gen_moe_cubin_pool/" - "trtllm_gen_moe_cubin_pool_20260617_v0613rc1\n" - "To use Marlin " - "instead, set --moe-runner-backend marlin explicitly." - ) - - if server_args.moe_runner_backend == "auto": - logger.info( - "Kimi-K3 on SM100/SM103: moe_runner_backend=flashinfer_mxfp4 " - "(trtllm-gen SiTU cubin pool found)." - ) - return {"moe_runner_backend": "flashinfer_mxfp4"} - return {} + logger.info( + "Kimi-K3 on SM100/SM103: moe_runner_backend=flashinfer_mxfp4 " + "(FlashInfer SiTU kernels)." + ) + return {"moe_runner_backend": "flashinfer_mxfp4"} @_register_for( diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 9a83da8a1..4e4ba29cf 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -1653,7 +1653,7 @@ def _set_envs_and_config(server_args: ServerArgs): if server_args.attention_backend == "flashinfer": assert_pkg_version( "flashinfer_python", - "0.6.15.post1", + "0.6.17", "Please uninstall the old version and " "reinstall the latest version by following the instructions " "at https://docs.flashinfer.ai/installation.html.", diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 45ddcd37a..cc0ab28b7 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -708,9 +708,6 @@ class Envs: # Launch the TRT-LLM MoE grouped GEMMs with PDL only at or below this # token count. SGLANG_TRTLLM_MOE_PDL_MAX_TOKENS = EnvInt(8192) - # Unpacked cubin pool for the JIT-built trtllm-gen fused MoE (cubins + flat - # ABI headers + overlay/). Unset means the path is unavailable, not empty. - SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL = EnvStr(None) # SGLang needs to know FlashInfer NVFP4 4over6 config to compute the global scale factor. FLASHINFER_NVFP4_4OVER6 = EnvBool(False) FLASHINFER_NVFP4_4OVER6_E4M3_USE_256 = EnvBool(False) @@ -1300,12 +1297,6 @@ class Envs: # ==================================================================== # Kimi-K3 - # TRT-LLM-gen fused MoE (SiTU) via sglang JIT: path to an unpacked SiTU - # cubin pool (cubins + flat ABI headers + overlay/; distributed as a - # single downloadable archive). Needs the public flashinfer package - # installed for the unmodified JIT sources. Unset = feature off. - SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL = EnvStr(None) - # MNNVL fused all-reduce (bf16, TP8): zero-copy 1shot multicast-push for # small messages and in-place NVLS 2shot on symmetric-memory tensors for # large ones, with an optional fused residual add. Covers the KDA o_proj diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 34c14f21f..28b0d2205 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -281,6 +281,7 @@ def fast_prefill_plan( fixed_split_size if fixed_split_size is not None else -1, False, # disable_split_kv 0, # num_colocated_ctas + 0, # uniform_q_len ] self._plan_info = self._cached_module.plan(*args) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index d77605385..4ad7bfc39 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -79,9 +79,11 @@ if is_flashinfer_available(): nvfp4_block_scale_interleave, trtllm_fp4_block_scale_moe, ) - from flashinfer.fused_moe.core import ( - get_w2_permute_indices_with_cache, + from flashinfer.fused_moe import ( + trtllm_fp4_block_scale_routed_moe, ) + from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache + from flashinfer.tllm_enums import ActivationType, RoutingMethodType # SM90 mixed-input helpers landed in FlashInfer #3084 (post-0.6.10). Older # versions don't ship them; gate at import so unrelated code paths still load. @@ -1531,20 +1533,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) if self.moe_runner_config.activation == "situ": - # SiTU is only in the private trtllm-gen cubin pool (the - # public artifact bakes swiglu into the fused-act cubins and - # silently computes the wrong activation). Routing must also - # be noaux_tc (sigmoid + correction bias, DeepSeekV3 method), - # not the renormalize-softmax default below. - from sglang.kernels.ops.moe import trtllm_gen_moe as situ_moe - - if not situ_moe.available(): - raise RuntimeError( - "activation='situ' with the flashinfer_mxfp4 runner " - "needs the SiTU cubin pool: set " - "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL (see " - "sglang/kernels/ops/moe/trtllm_gen_moe.py)." - ) + # FlashInfer 0.6.17+ ships the SiTU TRT-LLM-gen kernels. + # Routing must be noaux_tc (sigmoid + correction bias, + # DeepSeekV3), not the renormalize-softmax default below. # EP is cubin-internal: each rank computes its local expert slice # [offset, +num_local) and the caller all-reduces. ep=1 -> TP path. local_expert_offset = layer.moe_ep_rank * layer.num_local_experts @@ -1568,27 +1559,36 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) defer_finalize = _deferred_finalize_enabled.get() - result = situ_moe.trtllm_fp4_block_scale_routed_moe( - packed_topk_ids=packed_topk, + result = trtllm_fp4_block_scale_routed_moe( + topk_ids=packed_topk, + routing_bias=None, hidden_states=x_quant, hidden_states_scale=x_scale, gemm1_weights=layer.w13_weight, gemm1_weights_scale=layer.w13_weight_scale, + gemm1_bias=None, gemm1_alpha=layer.gemm1_alpha, - # SiTuGlu: gatedActBeta is the linear-half tanh - # clip; K3 stores it in gemm1_clamp_limit. + # SiTU beta is the linear-half tanh clip; K3 stores it + # in gemm1_clamp_limit. gemm1_beta=layer.gemm1_clamp_limit, + gemm1_clamp_limit=None, gemm2_weights=layer.w2_weight, gemm2_weights_scale=layer.w2_weight_scale, + gemm2_bias=None, output1_scale_scalar=None, output1_scale_gate_scalar=None, output2_scale_scalar=None, num_experts=layer.num_experts, top_k=packed_topk.shape[1], + n_group=None, + topk_group=None, intermediate_size=self.intermediate_size_per_partition, - activation_type=situ_moe.ACTIVATION_SITU, local_expert_offset=local_expert_offset, local_num_experts=layer.num_local_experts, + routed_scaling_factor=None, + routing_method_type=RoutingMethodType.TopK.value, + activation_type=ActivationType.Situ.value, + tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), output=symm_output, do_finalize=not defer_finalize, ) @@ -1604,6 +1604,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): expanded_idx_to_permuted_idx=expanded_idx, top_k=packed_topk.shape[1], ) + else: + result = result[0] return StandardCombineInput(hidden_states=result) # Bypassed topk: route from logits inside the op. @@ -1612,7 +1614,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): if bias_bf16 is None and correction_bias is not None: bias_bf16 = correction_bias.to(torch.bfloat16) layer._situ_routing_bias_bf16 = bias_bf16 - situ_moe.trtllm_fp4_block_scale_moe( + trtllm_fp4_block_scale_moe( # router_logits is a row-strided slice of the K3 fused # front GEMM output; the FFI reads it as dense. routing_logits=router_logits.to(torch.bfloat16).contiguous(), @@ -1621,12 +1623,15 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): hidden_states_scale=x_scale, gemm1_weights=layer.w13_weight, gemm1_weights_scale=layer.w13_weight_scale, + gemm1_bias=None, gemm1_alpha=layer.gemm1_alpha, - # SiTuGlu: gatedActBeta is the linear-half tanh clip; - # K3 stores it in gemm1_clamp_limit (situ_linear_beta). + # SiTU beta is the linear-half tanh clip; K3 stores it in + # gemm1_clamp_limit (situ_linear_beta). gemm1_beta=layer.gemm1_clamp_limit, + gemm1_clamp_limit=None, gemm2_weights=layer.w2_weight, gemm2_weights_scale=layer.w2_weight_scale, + gemm2_bias=None, output1_scale_scalar=None, output1_scale_gate_scalar=None, output2_scale_scalar=None, @@ -1638,11 +1643,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): routed_scaling_factor=( topk_output.topk_config.routed_scaling_factor or 1.0 ), - routing_method_type=situ_moe.ROUTING_DEEPSEEK_V3, - activation_type=situ_moe.ACTIVATION_SITU, + routing_method_type=RoutingMethodType.DeepSeekV3.value, + activation_type=ActivationType.Situ.value, norm_topk_prob=topk_output.topk_config.renormalize, local_expert_offset=local_expert_offset, local_num_experts=layer.num_local_experts, + tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), output=symm_output, ) return StandardCombineInput(hidden_states=symm_output) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 749a7e96a..466f72ea5 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1999,7 +1999,7 @@ def check_pkg_version_at_least(pkg: str, min_version: str) -> bool: Args: pkg: Package name (distribution name, e.g., "flashinfer-python") - min_version: Minimum version required (e.g., "0.6.15.post1") + min_version: Minimum version required (e.g., "0.6.17") Returns: True if package is installed and version >= min_version, False otherwise diff --git a/scripts/ci/cuda/ci_install_kimi_k3.sh b/scripts/ci/cuda/ci_install_kimi_k3.sh index bb8d7e45a..f66d4af9e 100755 --- a/scripts/ci/cuda/ci_install_kimi_k3.sh +++ b/scripts/ci/cuda/ci_install_kimi_k3.sh @@ -1,85 +1,14 @@ #!/bin/bash -# Install the standard CUDA CI dependencies plus Kimi-K3's FlashInfer assets. +# Install the standard CUDA CI dependencies plus Kimi-K3's Transformers fix. set -euxo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" # Source (not bash) so the generic install's Python/venv selection remains -# active while patching the FlashInfer package it just installed. +# active while applying the Transformers compatibility fix. # shellcheck disable=SC1091 source "${SCRIPT_DIR}/ci_install_dependency.sh" "$@" -TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip" -TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192" -TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1" -export SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool" - -install_required_tools() { - local missing=() - local tool - for tool in patch unzip wget; do - command -v "${tool}" >/dev/null 2>&1 || missing+=("${tool}") - done - if [ ${#missing[@]} -eq 0 ]; then - return - fi - - apt-get update || true - apt-get install -y --no-install-recommends "${missing[@]}" -} - -cubin_pool_is_valid() { - [ -d "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" ] && - [ "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" -type f -name '*.cubin' | wc -l)" -eq 1696 ] && - [ -f "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}/flashinferMetaInfo.h" ] && - [ -d "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}/local" ] && - [ -d "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}/overlay/csrc" ] -} - -install_trtllm_gen_moe_cubin_pool() ( - if cubin_pool_is_valid; then - echo "Reusing validated TRT-LLM Gen MoE cubin pool" - return - fi - - local cubin_archive - local cubin_extract_dir - local extracted_pool - cubin_archive="$(mktemp /tmp/trtllm_gen_moe_cubin_pool.XXXXXX.zip)" - cubin_extract_dir="$(mktemp -d /tmp/trtllm_gen_moe_cubin_extract.XXXXXX)" - extracted_pool="${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" - trap 'rm -f "${cubin_archive}"; rm -rf "${cubin_extract_dir}"' EXIT - - wget --no-verbose --output-document="${cubin_archive}" \ - "${TRTLLM_GEN_MOE_CUBIN_URL}" - echo "${TRTLLM_GEN_MOE_CUBIN_SHA256} ${cubin_archive}" | \ - sha256sum --check --strict - - unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" - test "$(find "${extracted_pool}" -type f -name '*.cubin' | wc -l)" -eq 1696 - - rm -rf "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" - mv "${extracted_pool}" "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" - cubin_pool_is_valid -) - -apply_flashinfer_dcp_patch() { - local flashinfer_dcp_patch - local flashinfer_site_packages - flashinfer_dcp_patch="${REPO_ROOT}/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt" - flashinfer_site_packages="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')" - - sed '/^diff --git a\/tests\//,$d' "${flashinfer_dcp_patch}" | \ - patch --dry-run --batch --forward --strip=1 \ - --directory="${flashinfer_site_packages}" - sed '/^diff --git a\/tests\//,$d' "${flashinfer_dcp_patch}" | \ - patch --batch --forward --strip=1 \ - --directory="${flashinfer_site_packages}" - - rm -rf /root/.cache/flashinfer /root/.cache/pip - python3 -c 'import inspect; from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla; assert "enable_dcp" in inspect.signature(trtllm_batch_decode_with_kv_cache_mla).parameters' -} - apply_transformers_symlink_patch() { # transformers 5.12.1 resolves custom-code symlinks out of the HF snapshot # and into blobs/, then looks for relative imports by filename in blobs/. @@ -126,14 +55,4 @@ with tempfile.TemporaryDirectory() as tmp: PY } -install_required_tools -install_trtllm_gen_moe_cubin_pool -apply_flashinfer_dcp_patch apply_transformers_symlink_patch - -# The install runs in its own shell. Persist the pool path for later workflow -# steps so Kimi-K3 selects the FlashInfer MXFP4 MoE runner instead of failing -# its startup validation. -if [ -n "${GITHUB_ENV:-}" ]; then - echo "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" >> "${GITHUB_ENV}" -fi diff --git a/test/registered/sampling/test_penalty.py b/test/registered/sampling/test_penalty.py index 63f7bc70a..5133d1fab 100644 --- a/test/registered/sampling/test_penalty.py +++ b/test/registered/sampling/test_penalty.py @@ -216,6 +216,9 @@ class TestPenalty(CustomTestCase): } self._test_penalty_effect(prompt, baseline_params, penalty_params) + @unittest.skip( + "TODO: Fix the flaky negative-penalty diversity assertion and re-enable." + ) def test_penalty_edge_cases_negative_penalty_values(self): """Test that negative penalties decrease vocabulary diversity.""" prompt = "Write the word 'test' exactly 15 times in a row, separated by spaces." diff --git a/test/registered/unit/mem_cache/test_unified_mamba_views.py b/test/registered/unit/mem_cache/test_unified_mamba_views.py index b7170b06b..b07988a78 100644 --- a/test/registered/unit/mem_cache/test_unified_mamba_views.py +++ b/test/registered/unit/mem_cache/test_unified_mamba_views.py @@ -313,7 +313,7 @@ def _k3_kda_mamba_geometry(heads_per_rank: int) -> dict: class TestKDAFlashInferEnvelopeStateContract(unittest.TestCase): """Derived property: the envelope-strided KDA temporal view (unified memory / page-major layout) must satisfy the state contract of FlashInfer - ``recurrent_kda`` (pinned ``flashinfer_python==0.6.14``), because the KDA + ``recurrent_kda`` (pinned ``flashinfer_python==0.6.17``), because the KDA flashinfer decode wrapper (``linear/kernels/kda_flashinfer.py``) passes the committed per-layer pool view straight into the kernel (in-place state update on the cu_seqlens path — no gather/scatter copy around the call).