From 5cb027f45e086242f9cded4175e50cd2ee6ec9a2 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sun, 18 Oct 2026 21:58:40 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../muse-glimmer/export/export_dflash.py | 145 +++++- .../models/muse-glimmer/export/export_solo.py | 90 +++- .../models/muse-glimmer/model/dflash_model.py | 59 +++ .../source_transformations/cuda.py | 101 +++- .../tests/test_cuda_dflash_packed.py | 488 ++++++++++++++++++ 5 files changed, 860 insertions(+), 23 deletions(-) create mode 100644 examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py diff --git a/examples/models/muse-glimmer/export/export_dflash.py b/examples/models/muse-glimmer/export/export_dflash.py index a8f82b65104..b4478bfccb1 100644 --- a/examples/models/muse-glimmer/export/export_dflash.py +++ b/examples/models/muse-glimmer/export/export_dflash.py @@ -24,8 +24,6 @@ def validate_dflash_export_options(backend: str, use_offgraph_kv_cache=False) -> None: - if use_offgraph_kv_cache and backend != "mlx": - raise ValueError("Off-graph DFlash requires MLX") if backend not in {"cuda", "mlx"}: raise ValueError(f"Unsupported DFlash backend: {backend}") @@ -129,6 +127,11 @@ def export_dflash( mmproj: str | None = None, max_vision_patches: int = 16384, use_offgraph_kv_cache: bool = False, + max_step_tokens: int = 4096, + max_cells: int = 65536, + draft_max_cells: int | None = None, + max_draft_owners: int = 4, + forward_sizes=None, ) -> None: """Export DFlash target + draft to one CUDA or MLX .pte. @@ -140,6 +143,8 @@ def export_dflash( mmproj loader and vision export contract to add ``vision_encoder``. """ validate_dflash_export_options(backend, use_offgraph_kv_cache) + if use_offgraph_kv_cache and backend == "cuda" and mmproj is not None: + raise ValueError("the CUDA batching DFlash export is text-only") import executorch.extension.llm.export.gguf # noqa: F401 import executorch.extension.llm.export.int4 # noqa: F401 @@ -197,6 +202,27 @@ def export_dflash( activation_dtype=activation_dtype, ) + if use_offgraph_kv_cache and backend == "cuda": + from executorch.examples.models.muse_glimmer.export.export_solo import ( + DEFAULT_FORWARD_SIZES, + export_cuda_batching, + ) + + export_cuda_batching( + target_model, + target_config, + output_dir, + max_step_tokens=max_step_tokens, + max_cells=max_cells, + forward_sizes=( + DEFAULT_FORWARD_SIZES if forward_sizes is None else forward_sizes + ), + dflash_draft=draft_model, + max_draft_owners=max_draft_owners, + draft_max_cells=draft_max_cells, + ) + return + vision_model = None pos_embed_table = None if mmproj is not None: @@ -229,6 +255,111 @@ def export_dflash( ) +def prepare_cuda_packed_draft( + target_model, + target_config, + draft_model, + *, + max_step_tokens: int, + max_cells: int, + max_owners: int, +): + """Trace the packed draft and publish its independent cell geometry.""" + from executorch.examples.models.muse_glimmer.export.export_solo import cell_manifest + from executorch.examples.models.muse_glimmer.model.dflash_model import ( + DFlashPackedExportModel, + ) + from executorch.examples.models.muse_glimmer.source_transformations.cuda import ( + add_dflash_hidden_tapping, + enable_dflash_offgraph_kv_cache, + offgraph_kv_cache_geometry, + ) + from torch.export import Dim, export + + config = draft_model.config + block = config.block_size + if ( + type(block) is not int + or block < 2 + or type(max_owners) is not int + or not 1 <= max_owners <= 4 + ): + raise ValueError("packed DFlash requires a trained block >= 2 and 1..4 owners") + if type(max_step_tokens) is not int or max_step_tokens < 2: + raise ValueError("packed DFlash max_step_tokens must be an integer >= 2") + max_write = max_step_tokens + max_owners * block + if type(max_cells) is not int or max_cells <= max_write: + raise ValueError("draft max_cells must exceed context plus all owner blocks") + if config.dim != target_config.dim: + raise ValueError("draft and target hidden dimensions must agree") + if target_model.activation_dtype != torch.bfloat16: + raise ValueError("CUDA packed DFlash requires bfloat16 activations") + if not 0 <= config.mask_token_id < target_config.vocab_size: + raise ValueError("draft mask_token_id must be in the target vocabulary") + add_dflash_hidden_tapping(target_model, list(config.target_layers)) + from executorch.examples.models.gemma4_31b.cuda_packers import ( + convert_quantized_tensors_for_cuda, + ) + + convert_quantized_tensors_for_cuda(draft_model) + draft_model.to(torch.bfloat16) + manifest = cell_manifest( + enable_dflash_offgraph_kv_cache(draft_model, max_write), max_cells + ) + combined = DFlashPackedExportModel(target_model, draft_model) + owners = ( + Dim.STATIC if max_owners == 1 else Dim("draft_owners", min=1, max=max_owners) + ) + context = Dim("draft_new_context", min=0, max=max_step_tokens) + kv_width = ( + context + block + if max_owners == 1 + else Dim("draft_kv_width", min=block, max=max_write) + ) + sample_owners = min(2, max_owners) + sample_queries = sample_owners * block + with torch.no_grad(): + program = export( + combined, + ( + torch.zeros((sample_owners, block), dtype=torch.long), + torch.zeros( + (1, 2, config.n_target_layers * config.dim), dtype=torch.bfloat16 + ), + torch.arange(2 + sample_queries, dtype=torch.long), + torch.arange(2 + sample_queries, dtype=torch.long), + ), + dynamic_shapes=( + {0: owners} if max_owners > 1 else None, + {1: context}, + {0: kv_width}, + {0: kv_width}, + ), + strict=True, + ) + geometry = offgraph_kv_cache_geometry(draft_model) + metadata = { + key.replace("get_", "get_draft_", 1): value for key, value in geometry.items() + } + metadata.update( + { + "get_block_size": block, + "get_mask_token_id": config.mask_token_id, + "get_n_target_layers": config.n_target_layers, + "get_dflash_target_layers": torch.tensor( + config.target_layers, dtype=torch.int32 + ), + "get_draft_offgraph_kv_max_cells": max_cells, + "get_draft_max_step_tokens": max_write, + "get_draft_max_owners": max_owners, + "get_draft_dflash_block_size": block, + "get_dflash_target_cache_key": "muse_glimmer_target", + "get_dflash_draft_cache_key": "muse_glimmer_draft", + } + ) + return {"draft_forward": program}, manifest, metadata + + def _dflash_constant_methods( *, target_config, @@ -926,6 +1057,11 @@ def main() -> None: "float16 for MLX and bfloat16 for CUDA.", ) parser.add_argument("--use-offgraph-kv-cache", action="store_true") + parser.add_argument("--max-step-tokens", type=int, default=4096) + parser.add_argument("--max-cells", type=int, default=65536) + parser.add_argument("--draft-max-cells", type=int, default=None) + parser.add_argument("--max-draft-owners", type=int, default=4) + parser.add_argument("--forward-sizes", default="1,2,4,8,16,32,64") args = parser.parse_args() try: @@ -955,6 +1091,11 @@ def main() -> None: mmproj=args.mmproj, max_vision_patches=args.max_vision_patches, use_offgraph_kv_cache=args.use_offgraph_kv_cache, + max_step_tokens=args.max_step_tokens, + max_cells=args.max_cells, + draft_max_cells=args.draft_max_cells, + max_draft_owners=args.max_draft_owners, + forward_sizes=[int(n) for n in args.forward_sizes.split(",") if n], ) diff --git a/examples/models/muse-glimmer/export/export_solo.py b/examples/models/muse-glimmer/export/export_solo.py index 4ae4b4d8afe..d7a0015b798 100644 --- a/examples/models/muse-glimmer/export/export_solo.py +++ b/examples/models/muse-glimmer/export/export_solo.py @@ -268,9 +268,21 @@ def step_forward( The whole step runs through the decoder so every token's K/V reaches the cache; only the rows the batch samples from go through the LM head. """ - x = model._run_blocks(model.embed_text(tokens), input_pos) + x = model.embed_text(tokens) + taps = None + if hasattr(model, "_dflash_target_layer_ids"): + sliding_mask, global_mask = model._build_masks(input_pos) + tapped = [] + for i, layer in enumerate(model.layers): + if i in model._dflash_target_layer_ids: + tapped.append(x) + x = layer(x, input_pos, sliding_mask, global_mask) + taps = torch.cat(tapped, dim=-1) + else: + x = model._run_blocks(x, input_pos) x = model.output_norm(x[:, logits_to_keep, :]) - return model._soft_cap(model.lm_head(x))[0] + logits = model._soft_cap(model.lm_head(x))[0] + return logits if taps is None else (logits, taps) def cell_manifest(sequence_manifest: str, max_cells: int) -> str: @@ -324,7 +336,8 @@ def device_residency(forward_methods, device_sampling: bool): return device -def export_cuda_batching( +# Keep target and draft assembly under one shared-weight collector. +def export_cuda_batching( # noqa: C901 model: MuseGlimmerModel, config: MuseGlimmerConfig, output_dir: str, @@ -332,6 +345,10 @@ def export_cuda_batching( max_cells: int, forward_sizes=DEFAULT_FORWARD_SIZES, device_sampling: bool = True, + *, + dflash_draft: nn.Module | None = None, + max_draft_owners: int = 4, + draft_max_cells: int | None = None, ): """Exports the forward methods and writes the artifact. @@ -362,6 +379,20 @@ def export_cuda_batching( ) from torch.export import Dim, export + if dflash_draft is not None: + from executorch.examples.models.muse_glimmer.export.export_dflash import ( + prepare_cuda_packed_draft, + ) + + draft_programs, draft_manifest, draft_metadata = prepare_cuda_packed_draft( + model, + config, + dflash_draft, + max_step_tokens=max_step_tokens, + max_cells=max_cells if draft_max_cells is None else draft_max_cells, + max_owners=max_draft_owners, + ) + forward_sizes = tuple(sorted({int(n) for n in forward_sizes})) widest_static = forward_sizes[-1] if forward_sizes else 0 if not forward_sizes or forward_sizes[0] != 1: @@ -392,6 +423,9 @@ def export_cuda_batching( manifest = cell_manifest( enable_offgraph_kv_cache(model, max_step_tokens), max_cells ) + if dflash_draft is not None: + # Deferred checkpoints still need scalar/RoPE buffers, after KV removal. + materialize_runtime_buffers(model, dtype=torch.bfloat16) geometry = offgraph_kv_cache_geometry(model) width = Dim( @@ -399,9 +433,10 @@ def export_cuda_batching( ) rows = Dim("rows", min=_CUDA_MIN_PREFILL_CHUNK, max=max_step_tokens) programs = {} - with common.BoundMethodForward( - model, types.MethodType(step_forward, model) - ), torch.no_grad(): + with ( + common.BoundMethodForward(model, types.MethodType(step_forward, model)), + torch.no_grad(), + ): for n in forward_sizes: print(f"Exporting forward_{n} (T={n})...") programs[f"forward_{n}"] = export( @@ -431,27 +466,46 @@ def export_cuda_batching( gc.collect() torch.cuda.empty_cache() forward_methods = list(programs) + if dflash_draft is not None: + programs.update(draft_programs) + del draft_programs, dflash_draft + cache_methods = forward_methods + ( + ["draft_forward"] if "draft_forward" in programs else [] + ) if device_sampling: print(f"Exporting {SAMPLE_METHOD} and {ARGMAX_METHOD}...") programs.update(sampler_programs(config.vocab_size)) def partitioner(name: str) -> CudaPartitioner: - if name not in forward_methods: + if name not in cache_methods: return CudaPartitioner( [ CudaBackend.generate_method_name_compile_spec(name), CompileSpec("autotune_at_compile_time", b"OFF"), ] ) - return CudaPartitioner( - [ - CudaBackend.generate_method_name_compile_spec(name), - CompileSpec("low_memory_mode", b"ON"), - CompileSpec(OFFGRAPH_KV_COMPILE_SPEC, manifest.encode()), - CompileSpec(OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, step_width_spec(name)), - CompileSpec("autotune_at_compile_time", b"OFF"), - ] - ) + is_draft = name == "draft_forward" + specs = [ + CudaBackend.generate_method_name_compile_spec(name), + CompileSpec("low_memory_mode", b"ON"), + CompileSpec( + OFFGRAPH_KV_COMPILE_SPEC, + (draft_manifest if is_draft else manifest).encode(), + ), + CompileSpec( + OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, + b"2:0" if is_draft else step_width_spec(name), + ), + CompileSpec("autotune_at_compile_time", b"OFF"), + ] + if "draft_forward" in cache_methods: + specs.append( + CompileSpec( + "llm_cache_registry_key", + b"muse_glimmer_draft" if is_draft else b"muse_glimmer_target", + ) + ) + return CudaPartitioner(specs) constant_methods = _solo_constant_methods( config=config, @@ -467,6 +521,8 @@ def partitioner(name: str) -> CudaPartitioner: constant_methods["get_min_prefill_chunk"] = _CUDA_MIN_PREFILL_CHUNK constant_methods["use_sampling"] = False constant_methods[MAX_CELLS_METHOD] = max_cells + if "draft_forward" in programs: + constant_methods.update(draft_metadata) print(f"Lowering {', '.join(programs)} to ExecuTorch (CUDA)...") edge = to_edge_transform_and_lower( @@ -487,7 +543,7 @@ def partitioner(name: str) -> CudaPartitioner: extract_delegate_segments=True, do_quant_fusion_and_const_prop=True, memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), - propagate_device_config=device_residency(forward_methods, device_sampling), + propagate_device_config=device_residency(cache_methods, device_sampling), emit_mutable_buffer_names=True, ), ) diff --git a/examples/models/muse-glimmer/model/dflash_model.py b/examples/models/muse-glimmer/model/dflash_model.py index bb21fac6676..2eb0b67559e 100644 --- a/examples/models/muse-glimmer/model/dflash_model.py +++ b/examples/models/muse-glimmer/model/dflash_model.py @@ -395,6 +395,46 @@ def forward( return self.norm(x) +class DFlashPackedExportModel(nn.Module): + """Packed export with shared weights named exactly as in target forwards.""" + + def __init__(self, target: nn.Module, draft: nn.Module): + super().__init__() + self.embed_tokens = target.embed_tokens + self.lm_head = target.lm_head + self.draft = draft + self.activation_dtype = target.activation_dtype + + def forward(self, draft_tokens, target_hidden, positions, kv_order): + return _packed_draft_forward( + self.embed_tokens, + self.lm_head, + self.draft, + self.activation_dtype, + draft_tokens, + target_hidden, + positions, + kv_order, + ) + + +def _packed_draft_forward( + embed_tokens, + lm_head, + draft, + activation_dtype, + draft_tokens, + target_hidden, + positions, + kv_order, +): + # Owners and blocks are separate axes, as in the MLX packed draft ABI. + torch._check(draft_tokens.shape[1] == draft.config.block_size) + torch._check(target_hidden.shape[0] == 1) + noise_embedding = embed_tokens(draft_tokens.reshape(1, -1)).to(activation_dtype) + return lm_head(draft(noise_embedding, target_hidden, positions, kv_order)).float() + + class MuseGlimmerWithDFlash(nn.Module): """Wrapper combining target and draft models with shared weights. @@ -460,6 +500,25 @@ def draft_forward( return self.target.lm_head(h).float() + def draft_forward_packed( + self, + draft_tokens: torch.Tensor, + target_hidden: torch.Tensor, + positions: torch.Tensor, + kv_order: torch.Tensor, + ) -> torch.Tensor: + """Cell-cache ABI; fourth input is explicitly not valid_ctx_len.""" + return _packed_draft_forward( + self.target.embed_tokens, + self.target.lm_head, + self.draft, + self.target.activation_dtype, + draft_tokens, + target_hidden, + positions, + kv_order, + ) + def draft_prefill( self, draft_tokens: torch.Tensor, diff --git a/examples/models/muse-glimmer/source_transformations/cuda.py b/examples/models/muse-glimmer/source_transformations/cuda.py index 191533a1b3f..635dfdca84a 100644 --- a/examples/models/muse-glimmer/source_transformations/cuda.py +++ b/examples/models/muse-glimmer/source_transformations/cuda.py @@ -21,13 +21,11 @@ # Register the TurboQuant attention operator. import executorch.backends.cuda.triton.kernels.tq4_sdpa # noqa: F401 import executorch.extension.llm.cache.update_and_attend # noqa: F401 - import torch import torch.nn as nn import torch.nn.functional as F import triton import triton.language as tl - from executorch.examples.models.muse_glimmer.model.dflash_model import ( build_dflash_swa_mask, rotate_half, @@ -40,15 +38,25 @@ from executorch.examples.models.muse_glimmer.source_transformations.sampler import ( sample, ) - from executorch.extension.llm.export.model_metadata import write_cache_geometry from torch.library import triton_op, wrap_triton def add_dflash_hidden_tapping(model: nn.Module, target_layer_ids: list[int]) -> None: """Make a CUDA Muse Glimmer target return logits and DFlash tapped hidden states.""" + if ( + not target_layer_ids + or any( + type(i) is not int or not 0 <= i < len(model.layers) + for i in target_layer_ids + ) + or target_layer_ids != sorted(set(target_layer_ids)) + ): + raise ValueError( + "DFlash tap IDs must be nonempty, ascending, unique and in range" + ) target_set = set(target_layer_ids) - model._dflash_target_layer_ids = target_layer_ids + model._dflash_target_layer_ids = tuple(target_layer_ids) def _run_embeddings_with_hidden( self, x: torch.Tensor, input_pos: torch.Tensor @@ -371,6 +379,91 @@ def dflash_cuda_source_transformations(model: nn.Module) -> None: ) +def _packed_dflash_attention_forward(self, x, target_hidden, positions, kv_order): + batch, width, _ = x.shape + context = target_hidden.shape[1] + torch._check(positions.shape[0] == context + width) + torch._check(kv_order.shape[0] == context + width) + # Project each real stream separately, preserving CUDA BF16 GEMM arithmetic. + k = torch.cat([self.k_proj(target_hidden), self.k_proj(x)], dim=1) + v = torch.cat([self.v_proj(target_hidden), self.v_proj(x)], dim=1) + k = self.k_norm(k.view(batch, context + width, self.n_kv_heads, self.head_dim)) + v = v.view(batch, context + width, self.n_kv_heads, self.head_dim) + q = self.q_norm(self.q_proj(x).view(batch, width, self.n_heads, self.head_dim)) + phase = torch.outer(positions.float(), self.inv_freq) + cos, sin = phase.cos()[None, None], phase.sin()[None, None] + q = rotate_half(q.transpose(1, 2), cos[:, :, context:], sin[:, :, context:]) + k = rotate_half(k.transpose(1, 2), cos, sin).index_select(2, kv_order) + v = v.transpose(1, 2).index_select(2, kv_order) + y = torch.ops.kvcache.update_and_attend( + q, + k, + v, + positions.index_select(0, kv_order).reshape(-1, 1), + self.layer_idx, + self.attn_scale, + q.dtype, + ) + return self.o_proj(y.transpose(1, 2).contiguous().reshape(batch, width, -1)) + + +def _packed_dflash_layer_forward(self, x, ctx, positions, kv_order): + x = x + self.self_attn(self.input_layernorm(x), ctx, positions, kv_order) + return x + self.mlp(self.post_attention_layernorm(x)) + + +def _packed_dflash_model_forward( + self, noise_embedding, target_hidden, positions, kv_order +): + ctx = self.hidden_norm(self.fc(target_hidden)) + x = noise_embedding + for layer in self.layers: + x = layer(x, ctx, positions, kv_order) + return self.norm(x) + + +def enable_dflash_offgraph_kv_cache(model: nn.Module, max_write: int) -> str: + """Packed cell-cache draft, with block-only queries and real context K/V.""" + layers = [] + for i, layer in enumerate(model.layers): + attn = layer.self_attn + attn.layer_idx = i + attn.register_buffer( + "inv_freq", + 1.0 + / ( + attn.rope_theta + ** ( + torch.arange(0, attn.head_dim, 2, dtype=torch.float32) + / attn.head_dim + ) + ), + persistent=False, + ) + layers.append( + { + "layer_id": i, + "policy": "ring" if attn.is_sliding else "flat", + "window": attn.window_size if attn.is_sliding else 0, + } + ) + if hasattr(attn, "kv_cache"): + del attn.kv_cache + attn.forward = types.MethodType(_packed_dflash_attention_forward, attn) + layer.forward = types.MethodType(_packed_dflash_layer_forward, layer) + model.forward = types.MethodType(_packed_dflash_model_forward, model) + return json.dumps( + { + "version": 1, + "maximum_capacity": model.config.max_seq_len, + "max_write": max_write, + "layers": layers, + }, + sort_keys=True, + separators=(",", ":"), + ) + + def materialize_dflash_runtime_buffers( model: nn.Module, dtype: torch.dtype, device: str = "cpu" ) -> None: diff --git a/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py b/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py new file mode 100644 index 00000000000..dbeb804f7a8 --- /dev/null +++ b/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py @@ -0,0 +1,488 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Independent tiny draft math and packed export shape checks, without MG weights.""" + +import copy +import importlib.resources +import json +import unittest +import unittest.mock +from dataclasses import replace +from pathlib import Path + +import torch +from executorch.examples.models.muse_glimmer.export import common +from executorch.examples.models.muse_glimmer.export.export_dflash import ( + prepare_cuda_packed_draft, +) +from executorch.examples.models.muse_glimmer.export.export_solo import step_forward +from executorch.examples.models.muse_glimmer.model.dflash_model import ( + DFlashConfig, + DFlashDraftModel, + MuseGlimmerWithDFlash, +) +from executorch.examples.models.muse_glimmer.model.model import MuseGlimmerModel +from executorch.examples.models.muse_glimmer.source_transformations.cuda import ( + add_dflash_hidden_tapping, + enable_dflash_offgraph_kv_cache, +) +from executorch.examples.models.muse_glimmer.tests.test_pipeline import TINY_CONFIG +from executorch.extension.llm.cache.reference_cache import AttendSpec, MaskKind +from executorch.extension.llm.cache.update_and_attend import REGISTRY + + +def models(block=4, head_dim=8): + torch.manual_seed(123) + tc = replace(TINY_CONFIG, dim=32, n_layers=2, head_dim=head_dim, max_seq_len=128) + target = MuseGlimmerModel(tc).to(torch.bfloat16).eval() + target.activation_dtype = torch.bfloat16 + dc = DFlashConfig( + dim=32, + n_layers=2, + n_heads=4, + n_kv_heads=2, + head_dim=head_dim, + ffn_dim=64, + vocab_size=tc.vocab_size, + target_layers=[0, 1], + block_size=block, + mask_token_id=255, + max_seq_len=128, + ) + draft = DFlashDraftModel(dc, max_context_length=64).to(torch.bfloat16).eval() + return target, draft + + +class PackedOracleCache: + """Test-only grouped noncausal cache: retain context, discard every block.""" + + def __init__(self, contexts, block): + self.contexts = contexts + self.block = block + self.history = {} + self.calls = [] + + def update_and_fetch(self, layer, k, v, positions): + self.calls.append((k.shape[2], positions.clone())) + specs = [] + at = 0 + for owner, context in enumerate(self.contexts): + end = at + context + self.block + new_k, new_v = k[:, :, at:end], v[:, :, at:end] + old = self.history.get((layer, owner)) + hk = new_k if old is None else torch.cat([old[0], new_k], dim=2) + hv = new_v if old is None else torch.cat([old[1], new_v], dim=2) + self.history[layer, owner] = ( + hk[:, :, : -self.block].clone(), + hv[:, :, : -self.block].clone(), + ) + specs.append(AttendSpec(hk, hv, MaskKind.NONE, self.block)) + at = end + if at != k.shape[2]: + raise ValueError("oracle grouping does not cover K/V") + return specs + + +class PackedDraftMathTest(unittest.TestCase): + def test_uneven_context_nonidentity_order_and_zero_new_context(self): + target, draft = models() + refs = [copy.deepcopy(draft), copy.deepcopy(draft)] + enable_dflash_offgraph_kv_cache(draft, 32) + combined = MuseGlimmerWithDFlash(target, draft, target.config, draft.config) + cache = PackedOracleCache([3, 1], 4) + key = "tiny-dflash-packed" + REGISTRY.install(key, cache) + self.addCleanup(REGISTRY.uninstall, key) + starts = [0, 0] + for contexts in ([3, 1], [0, 0], [1, 2]): + cache.contexts = contexts + tokens = torch.randint(0, 32, (2, 4)) + hidden = torch.randn(1, sum(contexts), 64, dtype=torch.bfloat16) + ctx_pos, block_pos, order = [], [], [] + offset = 0 + expected = [] + for owner, count in enumerate(contexts): + ctx_pos.extend(range(starts[owner], starts[owner] + count)) + block_pos.extend( + range(starts[owner] + count, starts[owner] + count + 4) + ) + order.extend(range(offset, offset + count)) + order.extend( + range(sum(contexts) + owner * 4, sum(contexts) + (owner + 1) * 4) + ) + with torch.no_grad(): + emb = target.embed_tokens(tokens[owner : owner + 1]).to( + torch.bfloat16 + ) + h = refs[owner]( + emb, + hidden[:, offset : offset + count], + torch.tensor([starts[owner]]), + ) + expected.append(target.lm_head(h).float()) + offset += count + starts[owner] += count + pos = torch.tensor(ctx_pos + block_pos) + with torch.no_grad(), REGISTRY.active(key): + actual = combined.draft_forward_packed( + tokens, hidden, pos, torch.tensor(order) + ) + torch.testing.assert_close( + actual, torch.cat(expected, dim=1), atol=0.025, rtol=0.025 + ) + self.assertEqual(cache.calls[-1][0], sum(contexts) + 8) + torch.testing.assert_close(cache.calls[-1][1][:, 0], pos[order]) + self.assertEqual(tuple(actual.shape), (1, 8, target.config.vocab_size)) + + def test_taps_keep_every_row_and_selected_logits_match(self): + target, _ = models() + ordinary = copy.deepcopy(target) + add_dflash_hidden_tapping(target, [0, 1]) + captured = [] + handles = [ + layer.register_forward_pre_hook( + lambda module, args: captured.append(args[0].clone()) + ) + for layer in ordinary.layers + ] + try: + tokens = torch.tensor([[1, 2, 3, 4, 5, 6]]) + positions = torch.arange(6) + with torch.no_grad(): + full = ordinary(tokens, positions) + logits, taps = step_forward( + target, tokens, positions, torch.tensor([5, 1, 5]) + ) + torch.testing.assert_close(logits, full[0, [5, 1, 5]]) + torch.testing.assert_close(taps, torch.cat(captured, dim=-1)) + self.assertEqual(tuple(taps.shape), (1, 6, 64)) + finally: + for handle in handles: + handle.remove() + + def test_invalid_tap_ids_fail_before_binding(self): + for ids in ([], [1, 0], [0, 0], [-1], [2], [True]): + target, _ = models() + with self.subTest(ids=ids), self.assertRaises(ValueError): + add_dflash_hidden_tapping(target, ids) + self.assertFalse(hasattr(target, "_dflash_target_layer_ids")) + + def test_invalid_packed_limits_fail_before_transform(self): + for overrides in ( + {"max_owners": True}, + {"max_owners": 0}, + {"max_owners": 5}, + {"max_step_tokens": 1}, + {"max_cells": 80}, + ): + target, draft = models(block=16) + options = {"max_step_tokens": 16, "max_cells": 128, "max_owners": 4} + options.update(overrides) + with self.subTest(overrides=overrides), self.assertRaises(ValueError): + prepare_cuda_packed_draft(target, target.config, draft, **options) + self.assertFalse(hasattr(target, "_dflash_target_layer_ids")) + self.assertTrue(hasattr(draft.layers[0].self_attn, "kv_cache")) + + def test_single_owner_export_and_malformed_width(self): + target, draft = models(block=16) + programs, _, metadata = prepare_cuda_packed_draft( + target, + target.config, + draft, + max_step_tokens=16, + max_cells=128, + max_owners=1, + ) + self.assertEqual(metadata["get_draft_max_step_tokens"], 32) + module = programs["draft_forward"].module() + key = "tiny-single-owner" + REGISTRY.install(key, PackedOracleCache([0], 16)) + self.addCleanup(REGISTRY.uninstall, key) + args = ( + torch.zeros(1, 16, dtype=torch.long), + torch.zeros(1, 0, 64, dtype=torch.bfloat16), + torch.arange(16), + torch.arange(16), + ) + with torch.no_grad(), REGISTRY.active(key): + self.assertEqual(tuple(module(*args).shape), (1, 16, 256)) + with self.assertRaises((RuntimeError, AssertionError)): + module(args[0], args[1], torch.arange(17), torch.arange(17)) + + def test_dynamic_export_matches_static_at_zero_and_unequal_context(self): + target, draft = models(block=16) + programs, manifest, metadata = prepare_cuda_packed_draft( + target, + target.config, + draft, + max_step_tokens=16, + max_cells=128, + max_owners=4, + ) + self.assertEqual(json.loads(manifest)["max_write"], 80) + self.assertEqual(metadata["get_block_size"], 16) + self.assertEqual(metadata["get_draft_max_owners"], 4) + self.assertEqual(metadata["get_draft_max_step_tokens"], 80) + state = programs["draft_forward"].state_dict + self.assertIs(state["embed_tokens.weight"], target.embed_tokens.weight) + self.assertIs(state["lm_head.weight"], target.lm_head.weight) + self.assertFalse(any(name.startswith("target.") for name in state)) + self.assertEqual(metadata["get_draft_n_caches"], 2) + self.assertNotEqual( + metadata["get_dflash_target_cache_key"], + metadata["get_dflash_draft_cache_key"], + ) + combined = MuseGlimmerWithDFlash(target, draft, target.config, draft.config) + key = "tiny-dynamic-draft" + for owners, context in ((1, 0), (1, 1), (2, 3), (3, 7), (4, 1), (4, 16)): + q = owners * 16 + args = ( + torch.zeros(owners, 16, dtype=torch.long), + torch.randn(1, context, 64, dtype=torch.bfloat16), + torch.arange(context + q), + torch.arange(context + q), + ) + with ( + common.BoundMethodForward(combined, combined.draft_forward_packed), + torch.no_grad(), + ): + static = torch.export.export(combined, args, strict=True) + counts = [context] + [0] * (owners - 1) + results = [] + for module in (programs["draft_forward"].module(), static.module()): + REGISTRY.install(key, PackedOracleCache(counts, 16)) + try: + with REGISTRY.active(key), torch.no_grad(): + results.append(module(*args)) + finally: + REGISTRY.uninstall(key) + torch.testing.assert_close(*results) + + +class PackedDraftCudaTest(unittest.TestCase): + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") + def test_exported_lowered_model_numerics_all_owner_counts(self): + from executorch.backends.cuda.passes.lower_offgraph_kv import ( + CheckOffGraphKVStepWidthPass, + LowerOffGraphKVPass, + OFFGRAPH_KV_CELLS_FQN, + offgraph_kv_mask_fqn, + OFFGRAPH_KV_READ_LEN_FQN, + parse_offgraph_kv_manifest, + ) + from executorch.exir import EdgeCompileConfig, to_edge + + target, draft = models(block=16, head_dim=16) + programs, manifest, _ = prepare_cuda_packed_draft( + target, + target.config, + draft, + max_step_tokens=16, + max_cells=128, + max_owners=4, + ) + edge = to_edge( + programs["draft_forward"], + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).exported_program() + CheckOffGraphKVStepWidthPass(b"2:0", layout="cell")(edge) + lowered = LowerOffGraphKVPass(parse_offgraph_kv_manifest(manifest.encode()))( + edge + ) + from executorch.backends.cuda.cuda_backend import _move_graph_nodes_to_device + + # Lowering deliberately frees runtime-owned cache storage for compilation. + # The numerical harness supplies actual storage before the device move. + for name, value in list(lowered.constants.items()): + if ( + isinstance(value, torch.Tensor) + and value.untyped_storage().nbytes() == 0 + ): + lowered.constants[name] = torch.zeros(value.shape, dtype=value.dtype) + # The lowered graph is intentionally mutable, so use CUDA's rewrite + # rather than the upstream functional-ATen export verifier. + _move_graph_nodes_to_device(lowered.graph_module, torch.device("cuda")) + oracle = MuseGlimmerWithDFlash( + target, draft, target.config, draft.config + ).cuda() + key = "tiny-lowered-model" + for counts in ( + [0], + [3, 1], + [0, 2, 4], + [0, 0, 0, 0], + [1, 0, 3, 2], + [4, 4, 4, 4], + ): + context, queries = sum(counts), 16 * len(counts) + order, ctx_pos, block_pos = [], [], [] + offset = 0 + mask = torch.zeros(1, 1, 80, 128, dtype=torch.bool, device="cuda") + grouped = 0 + for owner, count in enumerate(counts): + ctx_pos.extend(range(count)) + block_pos.extend(range(count, count + 16)) + order.extend(range(offset, offset + count)) + order.extend(range(context + owner * 16, context + (owner + 1) * 16)) + mask[ + :, :, owner * 16 : (owner + 1) * 16, grouped : grouped + count + 16 + ] = True + offset += count + grouped += count + 16 + args = ( + torch.randint(0, 256, (len(counts), 16), device="cuda"), + torch.randn(1, context, 64, dtype=torch.bfloat16, device="cuda"), + torch.tensor(ctx_pos + block_pos, device="cuda"), + torch.tensor(order, device="cuda"), + ) + lowered.constants[OFFGRAPH_KV_CELLS_FQN] = torch.arange(80, device="cuda") + lowered.constants[OFFGRAPH_KV_READ_LEN_FQN] = torch.tensor( + [grouped], device="cuda" + ) + for layer in range(2): + lowered.constants[f"__et_offgraph_kv_layer_{layer}_k"] = torch.zeros( + 1, 128, 2, 16, dtype=torch.bfloat16, device="cuda" + ) + lowered.constants[f"__et_offgraph_kv_layer_{layer}_v"] = torch.zeros( + 1, 128, 2, 16, dtype=torch.bfloat16, device="cuda" + ) + lowered.constants[offgraph_kv_mask_fqn(0)] = mask + module = lowered.module().cuda() + REGISTRY.install(key, PackedOracleCache(counts, 16)) + try: + with torch.no_grad(), REGISTRY.active(key): + expected = oracle.draft_forward_packed(*args) + actual = module(*args) + torch.testing.assert_close(actual, expected, atol=0.035, rtol=0.035) + self.assertEqual(tuple(actual.shape), (1, queries, 256)) + finally: + REGISTRY.uninstall(key) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") + def test_cell_step_writes_k_width_and_reads_only_block_queries(self): + from executorch.backends.cuda.passes.lower_offgraph_kv import cell_step + + torch.manual_seed(34) + q = torch.randn(1, 4, 8, 16, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, 2, 12, 16, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + cells = torch.tensor([11, 3, 8, 0, 9, 2, 4, 10, 1, 7, 5, 6], device="cuda") + mask = torch.zeros(1, 1, 16, 20, device="cuda", dtype=torch.bool) + mask[:, :, :4, cells[:7]] = True + mask[:, :, 4:8, cells[7:]] = True + kp = torch.zeros(1, 20, 2, 16, device="cuda", dtype=torch.bfloat16) + vp = torch.zeros_like(kp) + actual = cell_step( + q, k, v, kp, vp, cells, torch.tensor([12], device="cuda"), mask, 16**-0.5 + ) + expected = [] + for qs, ks in ((slice(0, 4), slice(0, 7)), (slice(4, 8), slice(7, 12))): + expected.append( + torch.nn.functional.scaled_dot_product_attention( + q[:, :, qs].float(), + k[:, :, ks].float().repeat_interleave(2, 1), + v[:, :, ks].float().repeat_interleave(2, 1), + scale=16**-0.5, + ).to(torch.bfloat16) + ) + torch.testing.assert_close(actual, torch.cat(expected, 2), atol=0.01, rtol=0.01) + torch.testing.assert_close(kp[:, cells].transpose(1, 2), k) + torch.testing.assert_close(vp[:, cells].transpose(1, 2), v) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") + def test_tiny_combined_lowering_schema(self): + import tempfile + + from executorch.examples.models.muse_glimmer.export.export_solo import ( + export_cuda_batching, + ) + from executorch.exir.lowered_backend_module import get_lowered_submodules + + target, draft = models(block=16, head_dim=16) + source = Path(__file__).resolve().parents[4] + read_resource = importlib.resources.read_binary + + def owned_schema(package, name): + # Source-only launchers omit the wheel's copied schema resources. + if package == "executorch.exir._serialize" and name in ( + "program.fbs", + "scalar_type.fbs", + ): + return (source / "schema" / name).read_bytes() + return read_resource(package, name) + + with ( + tempfile.TemporaryDirectory() as out, + unittest.mock.patch( + "importlib.resources.read_binary", side_effect=owned_schema + ), + ): + program = export_cuda_batching( + target, + target.config, + out, + max_step_tokens=16, + max_cells=128, + forward_sizes=(1, 2, 4, 8), + dflash_draft=draft, + max_draft_owners=4, + ) + self.assertGreater((Path(out) / "model.pte").stat().st_size, 0) + self.assertTrue(list(Path(out).glob("*.ptd"))) + self.assertEqual( + set(program.methods), + { + "forward_1", + "forward_2", + "forward_4", + "forward_8", + "forward_others", + "draft_forward", + "sample", + "sample_argmax", + }, + ) + from executorch.backends.cuda.cuda_weight_collector import ( + decode_cuda_aoti_metadata, + ) + + shared_keys = [] + for name in ("forward_4", "forward_others", "draft_forward"): + lowered = get_lowered_submodules( + program.exported_program(name).graph_module + ) + self.assertEqual(len(lowered), 1) + specs = {s.key: s.value for s in lowered[0][1].compile_specs} + self.assertEqual( + specs["offgraph_kv_step_width"], + b"2:0" if name == "draft_forward" else b"1:0", + ) + self.assertEqual( + specs["llm_cache_registry_key"], + ( + b"muse_glimmer_draft" + if name == "draft_forward" + else b"muse_glimmer_target" + ), + ) + self.assertEqual(specs["low_memory_mode"], b"ON") + self.assertEqual(specs["autotune_at_compile_time"], b"OFF") + weights = decode_cuda_aoti_metadata(lowered[0][1].processed_bytes) + shared_keys.append( + { + entry.fqn: entry.storage_key + for entry in weights.entries + if entry.fqn in ("embed_tokens.weight", "lm_head.weight") + } + ) + self.assertEqual(set(shared_keys[0]), {"embed_tokens.weight", "lm_head.weight"}) + self.assertEqual(shared_keys[0], shared_keys[1]) + self.assertEqual(shared_keys[0], shared_keys[2]) + + +if __name__ == "__main__": + unittest.main()