From 2c7b0770f1519e3ad7df1a4b67f98884915aba78 Mon Sep 17 00:00:00 2001 From: Moustafa Saleh <8815169+mesaleh@users.noreply.github.com> Date: Tue, 26 May 2026 09:43:34 -0700 Subject: [PATCH 1/4] feat(spec-decode): add native DFlash support Signed-off-by: Moustafa Saleh <8815169+mesaleh@users.noreply.github.com> --- docs/configuration/server.md | 2 +- docs/recipes/models.md | 32 ++ .../runtime/execution/cuda_graph_wrapper.py | 5 + .../runtime/execution/drafter/dflash.py | 506 ++++++++++++++++++ .../runtime/execution/model_executor.py | 33 +- .../runtime/execution/model_runner.py | 3 + .../runtime/layers/attention/backends/mha.py | 5 +- .../runtime/layers/attention/configs/mha.py | 6 +- .../tokenspeed/runtime/models/deepseek_v3.py | 16 + python/tokenspeed/runtime/models/dflash.py | 435 +++++++++++++++ .../runtime/spec_decode/algorithm.py | 2 + .../runtime/utils/hf_transformers_utils.py | 1 + .../tokenspeed/runtime/utils/server_args.py | 22 +- 13 files changed, 1058 insertions(+), 10 deletions(-) create mode 100644 python/tokenspeed/runtime/execution/drafter/dflash.py create mode 100644 python/tokenspeed/runtime/models/dflash.py diff --git a/docs/configuration/server.md b/docs/configuration/server.md index 905fbbd9a..091f3ce57 100644 --- a/docs/configuration/server.md +++ b/docs/configuration/server.md @@ -124,7 +124,7 @@ the values accepted by the bundled `tokenspeed-smg` package. | Parameter | Purpose | | --- | --- | | `--speculative-config` | JSON speculative decoding configuration. | -| `--speculative-algorithm` | Speculative algorithm, such as `EAGLE3` or `MTP`. | +| `--speculative-algorithm` | Speculative algorithm, such as `EAGLE3`, `MTP`, or `DFLASH`. | | `--speculative-draft-model-path` | Draft model path or repo ID. | | `--speculative-draft-model-quantization` | Draft model quantization. Defaults to `unquant`. | | `--speculative-num-steps` | Number of draft model steps. Defaults to `3`. | diff --git a/docs/recipes/models.md b/docs/recipes/models.md index a7ed90a31..2be2a886e 100644 --- a/docs/recipes/models.md +++ b/docs/recipes/models.md @@ -33,6 +33,38 @@ tokenspeed serve nvidia/Kimi-K2.5-NVFP4 \ For K2.6, keep the same parameter shape and change the checkpoint and parser only if the model card requires a different value. +To enable a compatible DFlash draft model, keep the target launch shape and add +the draft model path plus DFlash speculative decoding options: + +```bash +tokenspeed serve nvidia/Kimi-K2.6-NVFP4 \ + --served-model-name kimi-k2.6 \ + --trust-remote-code \ + --max-model-len 262144 \ + --kv-cache-dtype fp8 \ + --quantization nvfp4 \ + --tensor-parallel-size 4 \ + --enable-expert-parallel \ + --chunked-prefill-size 8192 \ + --max-num-seqs 256 \ + --attention-backend tokenspeed_mla \ + --moe-backend flashinfer_trtllm \ + --reasoning-parser kimi_k25 \ + --tool-call-parser kimik2 \ + --speculative-algorithm DFLASH \ + --speculative-draft-model-path /path/to/kimi-k2.6-dflash \ + --speculative-num-draft-tokens 8 \ + --speculative-num-steps 7 \ + --drafter-attention-backend fa4 \ + --host 0.0.0.0 \ + --port 8000 +``` + +Known limitation: native TokenSpeed DFlash currently uses full-history draft +attention. It does not yet expose an equivalent of SGLang's +`--speculative-dflash-draft-window-size`; add such a flag before relying on +bounded draft attention for long-context deployments. + ## GLM5 / GLM5.2 GLM5 launches usually need remote code, long context, expert parallelism, FP8 KV diff --git a/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py b/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py index 3941b0684..d290a07e0 100644 --- a/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py +++ b/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py @@ -430,6 +430,11 @@ def run_once(): torch.cuda.synchronize() dist.barrier() + # Warmups can switch a backend back to eager metadata objects. Restore + # the graph-backed metadata immediately before capture so replay-time + # metadata refreshes update the same tensors recorded by the graph. + self._init_capture_metadata(bs) + # Fill sampler buffers OUTSIDE the capture so RNG ops aren't recorded. if self.sampling_backend is not None: self.sampling_backend.prepare_capture( diff --git a/python/tokenspeed/runtime/execution/drafter/dflash.py b/python/tokenspeed/runtime/execution/drafter/dflash.py new file mode 100644 index 000000000..2d9b82a84 --- /dev/null +++ b/python/tokenspeed/runtime/execution/drafter/dflash.py @@ -0,0 +1,506 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.distributed.comm_ops import all_gather_into_tensor +from tokenspeed.runtime.execution.cache_loc_kernel import compute_out_cache_loc_uniform +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.execution.drafter.base import BaseDrafter +from tokenspeed.runtime.execution.forward_batch_info import ( + CaptureHiddenMode, + ForwardMode, +) +from tokenspeed.runtime.layers.logits_processor import LogitsMetadata +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.env import get_global_server_args +from tokenspeed.runtime.utils.nvtx import nvtx_range + +if TYPE_CHECKING: + from tokenspeed.runtime.execution.input_buffer import InputBuffers + from tokenspeed.runtime.execution.model_runner import ModelRunner + from tokenspeed.runtime.execution.runtime_states import RuntimeStates + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + +logger = get_colorful_logger(__name__) + + +class DFlash(BaseDrafter): + """DFlash block drafter backed by a native TokenSpeed draft model.""" + + def __init__( + self, + spec_num_tokens: int, + spec_num_steps: int, + page_size: int, + draft_model_runner: ModelRunner | None = None, + req_to_page: torch.Tensor | None = None, + attn_backend=None, + token_to_kv_pool=None, + runtime_states: RuntimeStates | None = None, + input_buffers: InputBuffers | None = None, + vocab_size: int | None = None, + ) -> None: + super().__init__( + spec_num_tokens=spec_num_tokens, + spec_num_steps=spec_num_steps, + draft_model_runner=draft_model_runner, + runtime_states=runtime_states, + input_buffers=input_buffers, + page_size=page_size, + req_to_page=req_to_page, + attn_backend=attn_backend, + token_to_kv_pool=token_to_kv_pool, + vocab_size=vocab_size, + ) + server_args = get_global_server_args() + if not server_args.speculative_draft_model_path: + raise ValueError("DFLASH requires --speculative-draft-model-path.") + + if draft_model_runner is None: + raise ValueError("Native DFLASH requires a draft model runner.") + self.device = torch.device(draft_model_runner.device) + self.model = draft_model_runner.model + + cfg = self.model.config + dflash_cfg = getattr(cfg, "dflash_config", {}) or {} + self.target_layer_ids = [int(x) for x in dflash_cfg.get("target_layer_ids", [])] + if not self.target_layer_ids: + raise ValueError( + "DFLASH draft config must define dflash_config.target_layer_ids." + ) + if "mask_token_id" not in dflash_cfg: + raise ValueError( + "DFLASH draft config must define dflash_config.mask_token_id." + ) + self.mask_token_id = int(dflash_cfg["mask_token_id"]) + self.block_size = int(getattr(cfg, "block_size", spec_num_tokens)) + if self.block_size != int(spec_num_tokens): + logger.warning( + "DFLASH block size mismatch: checkpoint block_size=%s, " + "runtime speculative_num_draft_tokens=%s.", + self.block_size, + spec_num_tokens, + ) + self.hidden_size = int(getattr(cfg, "hidden_size")) + self.idle_forward_steps = 1 + self._init_native_buffers() + self._greedy_gathered_max: torch.Tensor | None = None + self._greedy_gathered_ids: torch.Tensor | None = None + self._greedy_gather_cap = 0 + + def _init_native_buffers(self) -> None: + if self.input_buffers is None: + raise ValueError("Native DFLASH requires input buffers.") + if self.req_to_page is None: + raise ValueError("Native DFLASH requires req_to_page.") + if self.attn_backend is None or self.token_to_kv_pool is None: + raise ValueError("Native DFLASH requires draft attention components.") + + max_bs = self.input_buffers.max_bs + self.draft_seq_lens_buf = torch.zeros_like(self.input_buffers.seq_lens_buf) + self.draft_out_cache_loc_buf = torch.empty( + (max_bs * self.spec_num_tokens,), + dtype=torch.int32, + device=self.device, + ) + self.draft_input_lengths_buf = torch.full( + (max_bs,), + self.spec_num_tokens, + dtype=torch.int32, + device=self.device, + ) + self.draft_extend_seq_lens_cpu = torch.full( + (max_bs,), + self.spec_num_tokens, + dtype=torch.int32, + pin_memory=True, + ) + self.block_offsets = torch.arange( + self.spec_num_tokens, dtype=torch.int64, device=self.device + ) + self.block_ids_buf = torch.empty( + (max_bs, self.spec_num_tokens), dtype=torch.int32, device=self.device + ) + self.block_positions_buf = torch.empty( + (max_bs, self.spec_num_tokens), dtype=torch.int64, device=self.device + ) + + def bind_target_model(self, target_model) -> None: + language_model = getattr(target_model, "language_model", target_model) + self.target_model = target_model + self.target_language_model = language_model + self.embed_tokens = target_model.get_input_embeddings() + self.lm_head = target_model.lm_head + self.logits_processor = language_model.logits_processor + + def _greedy_sample_from_vocab_parallel_head( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + if not hasattr(self.lm_head, "weight") or not hasattr( + self.lm_head, "shard_indices" + ): + metadata = LogitsMetadata(forward_mode=ForwardMode.DECODE) + logits = self.logits_processor._get_logits( + hidden_states, self.lm_head, metadata + ) + return torch.argmax(logits, dim=-1).to(torch.int32) + + shard = self.lm_head.shard_indices + weight = self.lm_head.weight + hidden_states = hidden_states.to(weight.dtype) + + num_org = int(shard.num_org_elements) + num_org_padded = int(shard.num_org_elements_padded) + num_added = int(shard.num_added_elements) + org_vocab_start = int(shard.org_vocab_start_index) + added_vocab_start = int(shard.added_vocab_start_index) + + chunk_len = int(hidden_states.shape[0]) + if num_org > 0: + base_logits = torch.matmul(hidden_states, weight[:num_org].T) + local_max, local_arg = torch.max(base_logits, dim=-1) + else: + local_max = torch.full( + (chunk_len,), + torch.finfo(weight.dtype).min, + dtype=weight.dtype, + device=hidden_states.device, + ) + local_arg = torch.zeros( + (chunk_len,), dtype=torch.int64, device=hidden_states.device + ) + + if num_added > 0: + added_start = num_org_padded + added_end = num_org_padded + num_added + added_weight = weight[added_start:added_end] + added_logits = torch.matmul(hidden_states, added_weight.T) + added_max, added_arg = torch.max(added_logits, dim=-1) + use_added = added_max > local_max + local_max = torch.where(use_added, added_max, local_max) + local_arg = torch.where( + use_added, + added_arg.to(local_arg.dtype) + num_org_padded, + local_arg, + ) + + if num_added == 0: + global_ids = local_arg + org_vocab_start + else: + global_ids = torch.empty( + (chunk_len,), dtype=torch.int64, device=hidden_states.device + ) + is_base = local_arg < num_org + global_ids[is_base] = org_vocab_start + local_arg[is_base] + global_ids[~is_base] = added_vocab_start + ( + local_arg[~is_base] - num_org_padded + ) + + tp_size = int(self.logits_processor.tp_size) + if tp_size == 1: + return global_ids.to(torch.int32) + + needed = tp_size * chunk_len + if ( + self._greedy_gather_cap < needed + or self._greedy_gathered_max is None + or self._greedy_gathered_ids is None + or self._greedy_gathered_max.dtype != local_max.dtype + or self._greedy_gathered_max.device != hidden_states.device + ): + self._greedy_gathered_max = torch.empty( + (needed,), dtype=local_max.dtype, device=hidden_states.device + ) + self._greedy_gathered_ids = torch.empty( + (needed,), dtype=global_ids.dtype, device=hidden_states.device + ) + self._greedy_gather_cap = needed + + gathered_max = self._greedy_gathered_max[:needed] + gathered_ids = self._greedy_gathered_ids[:needed] + all_gather_into_tensor( + gathered_max, + local_max.contiguous(), + self.logits_processor.tp_rank, + self.logits_processor.tp_group, + ) + all_gather_into_tensor( + gathered_ids, + global_ids.contiguous(), + self.logits_processor.tp_rank, + self.logits_processor.tp_group, + ) + + gathered_max = gathered_max.view(tp_size, chunk_len) + gathered_ids = gathered_ids.view(tp_size, chunk_len) + best_rank = torch.argmax(gathered_max, dim=0).unsqueeze(0) + return torch.gather(gathered_ids, 0, best_rank).view(-1).to(torch.int32) + + @nvtx_range("dflash_update_native_cache", color="purple") + def _update_native_cache_from_target( + self, + base_ctx: ForwardContext, + logits_output: LogitsProcessorOutput, + accept_lengths: torch.Tensor, + ) -> None: + hidden = logits_output.hidden_states + if hidden is None: + raise RuntimeError("DFLASH requires target hidden states.") + if hidden.shape[0] != base_ctx.input_num_tokens: + raise RuntimeError( + "DFLASH hidden-state/token mismatch: " + f"hidden_tokens={hidden.shape[0]}, input_tokens={base_ctx.input_num_tokens}." + ) + + bs = base_ctx.bs + lengths = self.input_buffers.input_lengths_buf[:bs].to(torch.int64) + req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] + positions = self.input_buffers.positions_buf[: base_ctx.input_num_tokens] + cache_locs = self.input_buffers.out_cache_loc_buf[: base_ctx.input_num_tokens] + + if ( + base_ctx.num_extends == 0 + and torch.cuda.is_available() + and torch.cuda.is_current_stream_capturing() + ): + old_lens = self.runtime_states.valid_cache_lengths.index_select( + 0, req_pool_indices + ) + self.draft_seq_lens_buf[:bs].copy_( + old_lens.to(torch.int32) + accept_lengths[:bs].to(torch.int32) + ) + self._write_native_cache(hidden, positions, cache_locs) + return + + hidden_chunks = torch.split(hidden, lengths.detach().cpu().tolist(), dim=0) + pos_chunks = torch.split(positions, lengths.detach().cpu().tolist(), dim=0) + loc_chunks = torch.split(cache_locs, lengths.detach().cpu().tolist(), dim=0) + + selected_hidden = [] + selected_positions = [] + selected_cache_locs = [] + new_seq_lens = torch.empty((bs,), dtype=torch.int32, device=self.device) + + for row, (chunk, pos_chunk, loc_chunk) in enumerate( + zip(hidden_chunks, pos_chunks, loc_chunks, strict=True) + ): + if row < base_ctx.num_extends: + take = int(chunk.shape[0]) + else: + take = int(accept_lengths[row].item()) + if take <= 0: + pool_idx = req_pool_indices[row] + new_seq_lens[row] = self.runtime_states.valid_cache_lengths[pool_idx] + continue + + chunk = chunk[:take].contiguous() + pos_chunk = pos_chunk[:take].contiguous() + loc_chunk = loc_chunk[:take].contiguous() + selected_hidden.append(chunk) + selected_positions.append(pos_chunk) + selected_cache_locs.append(loc_chunk) + new_seq_lens[row] = (pos_chunk[-1] + 1).to(torch.int32) + + self.draft_seq_lens_buf[:bs].copy_(new_seq_lens) + if not selected_hidden: + return + + target_hidden = torch.cat(selected_hidden, dim=0) + target_positions = torch.cat(selected_positions, dim=0) + target_cache_locs = torch.cat(selected_cache_locs, dim=0) + self._write_native_cache(target_hidden, target_positions, target_cache_locs) + + def _write_native_cache( + self, + target_hidden: torch.Tensor, + target_positions: torch.Tensor, + target_cache_locs: torch.Tensor, + ) -> None: + target_hidden = target_hidden.to( + device=self.device, + dtype=self.draft_model_runner.model.fc.weight.dtype, + ) + expected_width = int(self.draft_model_runner.model.fc.in_features) + actual_width = int(target_hidden.shape[-1]) + if actual_width != expected_width: + raise RuntimeError( + "DFLASH captured hidden width mismatch: " + f"expected {expected_width}, got {actual_width}. " + "Check dflash_config.target_layer_ids against the target model." + ) + with torch.inference_mode(): + ctx_hidden = self.draft_model_runner.model.project_target_hidden( + target_hidden + ) + for layer in self.draft_model_runner.model.layers: + attn = layer.self_attn + k, v = attn.kv_proj_only(ctx_hidden) + k = attn.apply_k_norm(k) + k = attn.apply_k_rope(target_positions, k) + k = k.view(-1, attn.num_kv_heads, attn.head_dim) + v = v.view(-1, attn.num_kv_heads, attn.head_dim) + self.token_to_kv_pool.set_kv_buffer( + attn.attn, + target_cache_locs, + k, + v, + attn.attn.k_scale, + attn.attn.v_scale, + ) + + @staticmethod + def _current_tokens_from_output( + output_tokens: torch.Tensor, + accept_lengths: torch.Tensor, + num_extends: int, + spec_num_tokens: int, + ) -> torch.Tensor: + bs = accept_lengths.shape[0] + current = torch.empty((bs,), dtype=torch.int32, device=output_tokens.device) + if num_extends > 0: + current[:num_extends] = output_tokens[:num_extends] + num_decodes = bs - num_extends + if num_decodes > 0: + offsets = ( + torch.arange( + num_decodes, dtype=torch.int64, device=output_tokens.device + ) + * spec_num_tokens + - 1 + + num_extends + ) + current[num_extends:] = output_tokens[ + offsets + accept_lengths[num_extends:].to(torch.int64) + ] + return current + + def get_candidates(self, base_ctx: ForwardContext) -> torch.Tensor | None: + num_extends = base_ctx.num_extends + num_decodes = base_ctx.bs - num_extends + if num_decodes == 0: + return None + num_decode_tokens = num_decodes * self.spec_num_tokens + num_prefill_tokens = base_ctx.input_num_tokens - num_decode_tokens + return self.input_buffers.input_ids_buf[ + num_prefill_tokens : base_ctx.input_num_tokens + ].reshape(num_decodes, self.spec_num_tokens) + + def draft(self, current_tokens: torch.Tensor) -> torch.Tensor: + return self._draft_native(current_tokens) + + @nvtx_range("dflash_native_draft", color="purple") + def _draft_native(self, current_tokens: torch.Tensor) -> torch.Tensor: + bs = current_tokens.shape[0] + req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] + prefix_lens = self.draft_seq_lens_buf[:bs].clone() + seq_lens_after = self.draft_seq_lens_buf[:bs] + seq_lens_after.copy_(prefix_lens + int(self.spec_num_tokens)) + + block_ids = self.block_ids_buf[:bs] + block_ids.fill_(int(self.mask_token_id)) + block_ids[:, 0].copy_(current_tokens.to(torch.int32)) + block_positions = self.block_positions_buf[:bs] + block_positions.copy_( + prefix_lens.to(torch.int64).unsqueeze(1) + self.block_offsets + ) + + cache_locs = self.draft_out_cache_loc_buf[: bs * self.spec_num_tokens] + compute_out_cache_loc_uniform( + out_cache_loc_ptr=cache_locs, + req_pool_indices=req_pool_indices, + uniform_input_length=self.spec_num_tokens, + cache_start=prefix_lens, + req_to_pages=self.req_to_page, + page_size=self.page_size, + ) + + if not (torch.cuda.is_available() and torch.cuda.is_current_stream_capturing()): + self.attn_backend.init_forward_metadata( + bs=bs, + num_extends=0, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens_after, + req_to_page=self.req_to_page, + forward_mode=ForwardMode.DECODE, + extend_seq_lens_cpu=self.draft_extend_seq_lens_cpu[:bs], + ) + + ctx = ForwardContext( + attn_backend=self.attn_backend, + token_to_kv_pool=self.token_to_kv_pool, + req_to_page=self.req_to_page, + bs=bs, + num_extends=0, + input_num_tokens=bs * self.spec_num_tokens, + forward_mode=ForwardMode.DECODE, + capture_hidden_mode=CaptureHiddenMode.FULL, + ) + + flat_ids = block_ids.reshape(-1) + input_embeds = self.embed_tokens(flat_ids) + with torch.inference_mode(): + logits_output = self.draft_model_runner.forward( + ctx=ctx, + input_ids=flat_ids, + positions=block_positions.reshape(-1), + out_cache_loc=cache_locs, + input_lengths=self.draft_input_lengths_buf[:bs], + captured_hidden_states=None, + input_embeds=input_embeds, + ) + + draft_hidden = logits_output.hidden_states + if draft_hidden is None: + raise RuntimeError( + "Native DFLASH draft model did not return hidden states." + ) + draft_hidden = draft_hidden.view(bs, self.spec_num_tokens, self.hidden_size) + + next_tokens = torch.empty( + (bs, self.spec_num_tokens), dtype=torch.int32, device=self.device + ) + next_tokens[:, 0] = current_tokens.to(torch.int32) + sampled = self._greedy_sample_from_vocab_parallel_head( + draft_hidden[:, 1:, :].reshape(-1, self.hidden_size) + ) + next_tokens[:, 1:] = sampled.view(bs, self.spec_num_tokens - 1) + return next_tokens + + @nvtx_range("drafter:dflash", color="purple") + def run( + self, + base_ctx: ForwardContext, + logits_output: LogitsProcessorOutput, + output_tokens: torch.Tensor, + accept_lengths: torch.Tensor, + ) -> torch.Tensor: + if not hasattr(self, "target_model"): + raise RuntimeError("DFLASH drafter is not bound to a target model.") + self._update_native_cache_from_target(base_ctx, logits_output, accept_lengths) + current_tokens = self._current_tokens_from_output( + output_tokens, + accept_lengths, + base_ctx.num_extends, + self.spec_num_tokens, + ) + return self.draft(current_tokens) diff --git a/python/tokenspeed/runtime/execution/model_executor.py b/python/tokenspeed/runtime/execution/model_executor.py index 961490d91..e9a767ee3 100644 --- a/python/tokenspeed/runtime/execution/model_executor.py +++ b/python/tokenspeed/runtime/execution/model_executor.py @@ -36,6 +36,7 @@ from tokenspeed.runtime.execution.cache_loc_kernel import update_block_table from tokenspeed.runtime.execution.context import ForwardContext from tokenspeed.runtime.execution.cuda_graph_wrapper import CudaGraphWrapper +from tokenspeed.runtime.execution.drafter.dflash import DFlash from tokenspeed.runtime.execution.drafter.eagle import Eagle from tokenspeed.runtime.execution.forward_batch_info import ( CaptureHiddenMode, @@ -71,7 +72,7 @@ logger = get_colorful_logger(__name__) -_DRAFTER_MAPPING = {"EAGLE3": Eagle, "MTP": Eagle} +_DRAFTER_MAPPING = {"EAGLE3": Eagle, "MTP": Eagle, "DFLASH": DFlash} def _draft_idle_global_num_tokens_for_step( @@ -266,18 +267,34 @@ def __init__( token_to_kv_pool=draft_token_to_kv_pool, vocab_size=config.vocab_size, ) - embed, head = self.model_runner.model.get_embed_and_head() - draft_model_runner.model.set_embed_and_head(embed, head) + if hasattr(self.drafter, "bind_target_model"): + self.drafter.bind_target_model(self.model_runner.model) + # EAGLE3/MTP share the target's embed + lm_head; DFLASH ships its + # own draft weights, so it must NOT inherit the target's. + if config.spec_algo in ("EAGLE3", "MTP"): + embed, head = self.model_runner.model.get_embed_and_head() + draft_model_runner.model.set_embed_and_head(embed, head) target_hf = self.model_runner.model_config.hf_config mm_pad_substitute_id = getattr( target_hf, "image_token_id", None ) or getattr(target_hf, "media_placeholder_token_id", None) - if mm_pad_substitute_id is not None: + if mm_pad_substitute_id is not None and hasattr( + self.drafter, "set_mm_pad_substitute_id" + ): self.drafter.set_mm_pad_substitute_id(mm_pad_substitute_id) if config.spec_algo in ("EAGLE3",) and hasattr( self.model_runner.model, "set_eagle3_layers_to_capture" ): self.model_runner.model.set_eagle3_layers_to_capture() + if config.spec_algo == "DFLASH": + if not hasattr(self.model_runner.model, "set_dflash_layers_to_capture"): + raise ValueError( + "DFLASH requires the target model to support " + "set_dflash_layers_to_capture." + ) + self.model_runner.model.set_dflash_layers_to_capture( + self.drafter.target_layer_ids + ) else: self.drafter = None @@ -1076,7 +1093,13 @@ def execute_idle_forward( # NCCL collectives. Idle ranks must match those collectives: # 1 first-step forward + (spec_num_steps - 1) multi-step decode forwards. if self.drafter is not None: - for step_idx in range(self.drafter.spec_num_steps): + # DFLASH is a block drafter (idle_forward_steps=1); EAGLE3/MTP + # default to spec_num_steps. Mirror the active rank's per-step + # collective sizing either way. + idle_forward_steps = getattr( + self.drafter, "idle_forward_steps", self.drafter.spec_num_steps + ) + for step_idx in range(idle_forward_steps or 0): # Mirror active rank's catch-up step: when all non-idle ranks # are decoding, step 0 sizes collectives from bs/global_bs. draft_global_num_tokens = _draft_idle_global_num_tokens_for_step( diff --git a/python/tokenspeed/runtime/execution/model_runner.py b/python/tokenspeed/runtime/execution/model_runner.py index 29ce1c243..099928bb9 100644 --- a/python/tokenspeed/runtime/execution/model_runner.py +++ b/python/tokenspeed/runtime/execution/model_runner.py @@ -135,6 +135,7 @@ def forward( seq_lens: torch.Tensor | None = None, extend_prefix_lens: torch.Tensor | None = None, captured_hidden_states: torch.Tensor | None = None, + input_embeds: torch.Tensor | None = None, multimodal_context: MultimodalForwardContext | None = None, spec_step_idx: int | None = None, ) -> LogitsProcessorOutput: @@ -149,6 +150,8 @@ def forward( kwargs["get_embedding"] = True if captured_hidden_states is not None: kwargs["captured_hidden_states"] = captured_hidden_states + if input_embeds is not None: + kwargs["input_embeds"] = input_embeds if multimodal_context is not None: kwargs["multimodal_context"] = multimodal_context if spec_step_idx is not None and getattr( diff --git a/python/tokenspeed/runtime/layers/attention/backends/mha.py b/python/tokenspeed/runtime/layers/attention/backends/mha.py index c424d7398..757e553ff 100644 --- a/python/tokenspeed/runtime/layers/attention/backends/mha.py +++ b/python/tokenspeed/runtime/layers/attention/backends/mha.py @@ -426,7 +426,10 @@ def _forward_extend( cache_seqlens=metadata.seq_lens, max_seqlen_q=metadata.max_extend_seq_len, max_seqlen_k=self.max_context_len, - is_causal=True, + # DFLASH marks its draft attention non-causal so the draft block's + # query positions attend bidirectionally. Every other layer leaves + # the attribute unset, so this stays causal by default. + is_causal=not bool(getattr(layer, "non_causal", False)), window_left=layer.sliding_window_size, logit_cap=layer.logit_cap, sinks=sinks, diff --git a/python/tokenspeed/runtime/layers/attention/configs/mha.py b/python/tokenspeed/runtime/layers/attention/configs/mha.py index 22fcb7f42..37fad45fc 100644 --- a/python/tokenspeed/runtime/layers/attention/configs/mha.py +++ b/python/tokenspeed/runtime/layers/attention/configs/mha.py @@ -45,6 +45,10 @@ def generate( speculative_num_steps=server_args.speculative_num_steps, speculative_num_draft_tokens=server_args.speculative_num_draft_tokens, ) + kv_cache_dtype = server_args.kv_cache_dtype + if is_draft and server_args.speculative_algorithm == "DFLASH": + kv_cache_dtype = "bfloat16" + return cls( device=server_args.device, context_len=model_config.context_len, @@ -58,7 +62,7 @@ def generate( head_dim=model_config.head_dim, attn_tp_size=server_args.attn_tp_size or server_args.mapping.attn.tp_size, dtype=model_config.dtype, - kv_cache_dtype=resolve_dtype(server_args.kv_cache_dtype), + kv_cache_dtype=resolve_dtype(kv_cache_dtype), page_size=server_args.block_size, max_bs=server_args.max_num_seqs // (server_args.data_parallel_size or server_args.mapping.attn.dp_size), diff --git a/python/tokenspeed/runtime/models/deepseek_v3.py b/python/tokenspeed/runtime/models/deepseek_v3.py index ab963a810..0b737ae4a 100644 --- a/python/tokenspeed/runtime/models/deepseek_v3.py +++ b/python/tokenspeed/runtime/models/deepseek_v3.py @@ -1410,6 +1410,22 @@ def set_eagle3_layers_to_capture(self, layer_ids: list[int] | None = None): else: self.model.layers_to_capture = {val + 1 for val in layer_ids} + def set_dflash_layers_to_capture(self, layer_ids: list[int]): + # DFlash checkpoints name 0-indexed target layer outputs. The capture + # check runs before layer i, so capture at i + 1 for layer i's output. + num_layers = len(self.model.layers) + if len(set(layer_ids)) != len(layer_ids): + raise ValueError("DFLASH target_layer_ids must be unique.") + + invalid = [val for val in layer_ids if val < 0 or val + 1 >= num_layers] + if invalid: + raise ValueError( + "DFLASH target_layer_ids must map to capturable target layer " + f"outputs. Got invalid ids {invalid}; valid range is " + f"[0, {num_layers - 2}] for {num_layers} target layers." + ) + self.model.layers_to_capture = {val + 1 for val in layer_ids} + def get_param(self, params_dict, name): if name in params_dict: return params_dict[name] diff --git a/python/tokenspeed/runtime/models/dflash.py b/python/tokenspeed/runtime/models/dflash.py new file mode 100644 index 000000000..993dfc970 --- /dev/null +++ b/python/tokenspeed/runtime/models/dflash.py @@ -0,0 +1,435 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Optional + +import torch +from torch import nn + +from tokenspeed.runtime.distributed.comm_ops import all_reduce +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.activation import SiluAndMul +from tokenspeed.runtime.layers.layernorm import RMSNorm +from tokenspeed.runtime.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.utils import add_prefix +from tokenspeed.runtime.utils.env import global_server_args_dict + + +class DFlashAttention(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.mapping = mapping + self.hidden_size = int(config.hidden_size) + self.tp_rank = self.mapping.attn.tp_rank + self.tp_size = self.mapping.attn.tp_size + self.total_num_heads = int(config.num_attention_heads) + self.total_num_kv_heads = int( + getattr(config, "num_key_value_heads", self.total_num_heads) + ) + assert self.total_num_heads % self.tp_size == 0 + if self.total_num_kv_heads >= self.tp_size: + assert self.total_num_kv_heads % self.tp_size == 0 + else: + assert self.tp_size % self.total_num_kv_heads == 0 + + self.num_heads = self.total_num_heads // self.tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size) + self.head_dim = int( + getattr(config, "head_dim", self.hidden_size // self.total_num_heads) + ) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=bool(getattr(config, "attention_bias", False)), + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=bool(getattr(config, "attention_bias", False)), + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + reduce_results=False, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + eps = float(getattr(config, "rms_norm_eps", 1e-6)) + self.q_norm = RMSNorm(self.head_dim, eps=eps) + self.k_norm = RMSNorm(self.head_dim, eps=eps) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=int(getattr(config, "max_position_embeddings", 32768)), + base=float(getattr(config, "rope_theta", 1000000)), + rope_scaling=getattr(config, "rope_scaling", None), + ) + + # The FA4 MHA extend selector currently has no sliding-window kernel + # for this draft shape. Use full attention for draft proposals; target + # verification remains authoritative for accepted tokens. + sliding_window = -1 + self.attn = PagedAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + sliding_window_size=sliding_window, + ) + self.attn.non_causal = True + + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + q = q.reshape(-1, self.head_dim) + k = k.reshape(-1, self.head_dim) + q = self.q_norm(q).view(-1, self.q_size) + k = self.k_norm(k).view(-1, self.kv_size) + return q, k + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self._apply_qk_norm(q, k) + q, k = self.rotary_emb(positions, q, k) + k_cache = k.view(-1, self.num_kv_heads, self.head_dim) + v_cache = v.view(-1, self.num_kv_heads, self.head_dim) + ctx.token_to_kv_pool.set_kv_buffer( + self.attn, + out_cache_loc, + k_cache, + v_cache, + self.attn.k_scale, + self.attn.v_scale, + ) + attn_output = self.attn( + q, + None, + None, + ctx, + out_cache_loc, + save_kv_cache=False, + ) + if len(attn_output.size()) == 3: + attn_output = attn_output.reshape(attn_output.shape[0], -1) + output, _ = self.o_proj(attn_output) + return output + + def kv_proj_only( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + qkv, _ = self.qkv_proj(hidden_states) + _, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + return k, v + + def apply_k_norm(self, k: torch.Tensor) -> torch.Tensor: + k_shape = k.shape + return self.k_norm(k.reshape(-1, self.head_dim)).view(k_shape) + + def apply_k_rope(self, positions: torch.Tensor, k: torch.Tensor) -> torch.Tensor: + dummy_q = k.new_empty(k.shape) + _, k = self.rotary_emb(positions, dummy_q, k) + return k + + +class DFlashMLP(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + hidden_size = int(config.hidden_size) + intermediate_size = int(config.intermediate_size) + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + tp_rank=mapping.dense.tp_rank, + tp_size=mapping.dense.tp_size, + tp_group=mapping.dense.tp_group, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + reduce_results=False, + tp_rank=mapping.dense.tp_rank, + tp_size=mapping.dense.tp_size, + tp_group=mapping.dense.tp_group, + ) + if getattr(config, "hidden_act", "silu") != "silu": + raise ValueError("DFlash only supports silu activation.") + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class DFlashDecoderLayer(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + hidden_size = int(config.hidden_size) + eps = float(getattr(config, "rms_norm_eps", 1e-6)) + self.mapping = mapping + self.input_layernorm = RMSNorm(hidden_size, eps=eps) + self.self_attn = DFlashAttention( + config=config, + mapping=mapping, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.post_attention_layernorm = RMSNorm(hidden_size, eps=eps) + self.mlp = DFlashMLP( + config=config, + mapping=mapping, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: Optional[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + if ctx.forward_mode.is_idle(): + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + elif ( + ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"] + ): + hidden_states = all_reduce( + hidden_states, self.mapping.dense.tp_rank, self.mapping.dense.tp_group + ) + hidden_states, residual = self.input_layernorm(hidden_states, residual) + else: + hidden_states, residual, *_ = ( + self.input_layernorm.forward_with_allreduce_fusion( + self.mapping.dense.tp_rank, + self.mapping.dense.tp_group, + hidden_states, + residual, + ) + ) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + ) + + if ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"]: + hidden_states = all_reduce( + hidden_states, self.mapping.attn.tp_rank, self.mapping.attn.tp_group + ) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + else: + hidden_states, residual, *_ = ( + self.post_attention_layernorm.forward_with_allreduce_fusion( + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + hidden_states, + residual, + ) + ) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class DFlashDraftModel(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + eps = float(getattr(config, "rms_norm_eps", 1e-6)) + self.layers = nn.ModuleList( + [ + DFlashDecoderLayer( + config=config, + mapping=mapping, + layer_id=i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(int(config.num_hidden_layers)) + ] + ) + self.norm = RMSNorm(int(config.hidden_size), eps=eps) + target_layer_ids = (getattr(config, "dflash_config", {}) or {}).get( + "target_layer_ids", [] + ) + self.num_context_features = len(target_layer_ids) + self.fc = nn.Linear( + self.num_context_features * int(config.hidden_size), + int(config.hidden_size), + bias=False, + ) + self.hidden_norm = RMSNorm(int(config.hidden_size), eps=eps) + self.block_size = int(getattr(config, "block_size", 8)) + + def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor: + return self.hidden_norm(self.fc(target_hidden)) + + @torch.no_grad() + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + input_lengths: torch.Tensor, + input_embeds: torch.Tensor | None = None, + **kwargs, + ) -> LogitsProcessorOutput: + if input_embeds is None: + if not ctx.forward_mode.is_idle(): + raise ValueError("DFlashDraftModel requires input_embeds.") + hidden_states = self.fc.weight.new_empty((0, int(self.config.hidden_size))) + else: + hidden_states = input_embeds + residual = None + + for layer in self.layers: + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + residual=residual, + ) + + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + + return LogitsProcessorOutput( + next_token_logits=None, hidden_states=hidden_states + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + + def resolve_name(name: str) -> str | None: + if name in params_dict: + return name + if name.startswith("model.") and name[len("model.") :] in params_dict: + return name[len("model.") :] + return None + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if f".{weight_name}." not in name: + continue + resolved = resolve_name(name.replace(weight_name, param_name)) + if resolved is None: + continue + param = params_dict[resolved] + param.weight_loader(param, loaded_weight, shard_id) + break + else: + resolved = resolve_name(name) + if resolved is None: + continue + param = params_dict[resolved] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = [DFlashDraftModel] diff --git a/python/tokenspeed/runtime/spec_decode/algorithm.py b/python/tokenspeed/runtime/spec_decode/algorithm.py index 7f7b10ba5..dfedfddaa 100755 --- a/python/tokenspeed/runtime/spec_decode/algorithm.py +++ b/python/tokenspeed/runtime/spec_decode/algorithm.py @@ -25,6 +25,7 @@ class SpeculativeAlgorithm(IntEnum): NONE = auto() EAGLE3 = auto() MTP = auto() + DFLASH = auto() def is_none(self) -> bool: return self == SpeculativeAlgorithm.NONE @@ -37,6 +38,7 @@ def from_string(name: str | None) -> "SpeculativeAlgorithm": name_map = { "EAGLE3": SpeculativeAlgorithm.EAGLE3, "MTP": SpeculativeAlgorithm.MTP, + "DFLASH": SpeculativeAlgorithm.DFLASH, None: SpeculativeAlgorithm.NONE, } if name is not None: diff --git a/python/tokenspeed/runtime/utils/hf_transformers_utils.py b/python/tokenspeed/runtime/utils/hf_transformers_utils.py index 13797feee..4dbe43383 100755 --- a/python/tokenspeed/runtime/utils/hf_transformers_utils.py +++ b/python/tokenspeed/runtime/utils/hf_transformers_utils.py @@ -260,6 +260,7 @@ def get_config( and config.architectures and "NextN" not in config.architectures[0] and "Eagle" not in config.architectures[0] + and "DFlash" not in config.architectures[0] ): if config.architectures[0] == "MiniMaxM2ForCausalLM": config.architectures[0] = "LlamaForCausalLMEagle3" diff --git a/python/tokenspeed/runtime/utils/server_args.py b/python/tokenspeed/runtime/utils/server_args.py index 50102d38d..7410e2101 100755 --- a/python/tokenspeed/runtime/utils/server_args.py +++ b/python/tokenspeed/runtime/utils/server_args.py @@ -341,7 +341,13 @@ def resolve_config_aliases(self): num_speculative_tokens = config.get("num_speculative_tokens") if num_speculative_tokens is not None: - self.speculative_num_steps = int(num_speculative_tokens) + num_speculative_tokens = int(num_speculative_tokens) + if self.speculative_algorithm == "DFLASH": + if self.speculative_num_draft_tokens is None: + self.speculative_num_draft_tokens = num_speculative_tokens + self.speculative_num_steps = max(num_speculative_tokens - 1, 0) + else: + self.speculative_num_steps = num_speculative_tokens if self.speculative_num_draft_tokens is None: self.speculative_num_draft_tokens = self.speculative_num_steps + 1 @@ -521,6 +527,18 @@ def resolve_speculative_decoding(self): if self.speculative_draft_model_quantization == "unquant": self.speculative_draft_model_quantization = None + if self.speculative_algorithm == "DFLASH": + expected_steps = max(int(self.speculative_num_draft_tokens) - 1, 0) + if self.speculative_num_steps == ServerArgs.speculative_num_steps: + self.speculative_num_steps = expected_steps + elif self.speculative_num_steps != expected_steps: + raise ValueError( + "DFLASH requires speculative_num_steps to equal " + "speculative_num_draft_tokens - 1. " + f"Got {self.speculative_num_steps=} and " + f"{self.speculative_num_draft_tokens=}." + ) + if self.eagle3_layers_to_capture is not None: self.eagle3_layers_to_capture = [ int(x) for x in self.eagle3_layers_to_capture.split(",") @@ -1422,7 +1440,7 @@ def add_cli_args(parser: argparse.ArgumentParser): parser.add_argument( "--speculative-algorithm", type=str, - choices=["EAGLE3", "MTP"], + choices=["EAGLE3", "MTP", "DFLASH"], help="Speculative algorithm.", ) parser.add_argument( From 71dcfd84bdf67dbe2d019b5800c361ba18fa7105 Mon Sep 17 00:00:00 2001 From: Yue Weng <25103990+yweng0828@users.noreply.github.com> Date: Wed, 17 Jun 2026 20:50:40 -0700 Subject: [PATCH 2/4] fix some bugs and add cuda graph support --- .../runtime/execution/drafter/dflash.py | 29 +++- .../runtime/layers/attention/backends/mha.py | 143 ++++++++++++++++-- .../runtime/layers/attention/configs/base.py | 4 + .../runtime/layers/attention/configs/mha.py | 6 +- python/tokenspeed/runtime/models/dflash.py | 10 +- 5 files changed, 166 insertions(+), 26 deletions(-) diff --git a/python/tokenspeed/runtime/execution/drafter/dflash.py b/python/tokenspeed/runtime/execution/drafter/dflash.py index 2d9b82a84..fb2c54e06 100644 --- a/python/tokenspeed/runtime/execution/drafter/dflash.py +++ b/python/tokenspeed/runtime/execution/drafter/dflash.py @@ -31,7 +31,6 @@ ) from tokenspeed.runtime.layers.logits_processor import LogitsMetadata from tokenspeed.runtime.utils import get_colorful_logger -from tokenspeed.runtime.utils.env import get_global_server_args from tokenspeed.runtime.utils.nvtx import nvtx_range if TYPE_CHECKING: @@ -71,12 +70,13 @@ def __init__( token_to_kv_pool=token_to_kv_pool, vocab_size=vocab_size, ) - server_args = get_global_server_args() + if draft_model_runner is None: + raise ValueError("Native DFLASH requires a draft model runner.") + + server_args = draft_model_runner.server_args if not server_args.speculative_draft_model_path: raise ValueError("DFLASH requires --speculative-draft-model-path.") - if draft_model_runner is None: - raise ValueError("Native DFLASH requires a draft model runner.") self.device = torch.device(draft_model_runner.device) self.model = draft_model_runner.model @@ -241,13 +241,11 @@ def _greedy_sample_from_vocab_parallel_head( all_gather_into_tensor( gathered_max, local_max.contiguous(), - self.logits_processor.tp_rank, self.logits_processor.tp_group, ) all_gather_into_tensor( gathered_ids, global_ids.contiguous(), - self.logits_processor.tp_rank, self.logits_processor.tp_group, ) @@ -273,7 +271,12 @@ def _update_native_cache_from_target( ) bs = base_ctx.bs - lengths = self.input_buffers.input_lengths_buf[:bs].to(torch.int64) + # The target verify forward emits spec_num_tokens hidden states per + # decode request (the candidate block); input_lengths_buf only tracks + # the committed-token count there, so split decode rows by + # spec_num_tokens. Prefill rows keep their real chunk lengths. + lengths = self.input_buffers.input_lengths_buf[:bs].to(torch.int64).clone() + lengths[base_ctx.num_extends :] = self.spec_num_tokens req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] positions = self.input_buffers.positions_buf[: base_ctx.input_num_tokens] cache_locs = self.input_buffers.out_cache_loc_buf[: base_ctx.input_num_tokens] @@ -442,8 +445,19 @@ def _draft_native(self, current_tokens: torch.Tensor) -> torch.Tensor: seq_lens=seq_lens_after, req_to_page=self.req_to_page, forward_mode=ForwardMode.DECODE, + # Draft block runs in DECODE mode; the extend_* params are + # required by the signature but unused on the decode path. + extend_seq_lens=None, extend_seq_lens_cpu=self.draft_extend_seq_lens_cpu[:bs], + extend_prefix_lens=None, + extend_prefix_lens_cpu=None, ) + else: + # CUDA-graph capture/replay: the expanded decode metadata + # (page table) is prepared out-of-graph by the wrapper; broadcast + # the live per-request block-end length into the expanded seq_lens + # buffer here so the recorded op re-derives them on every replay. + self.attn_backend.fill_block_decode_seq_lens(bs, seq_lens_after) ctx = ForwardContext( attn_backend=self.attn_backend, @@ -464,7 +478,6 @@ def _draft_native(self, current_tokens: torch.Tensor) -> torch.Tensor: input_ids=flat_ids, positions=block_positions.reshape(-1), out_cache_loc=cache_locs, - input_lengths=self.draft_input_lengths_buf[:bs], captured_hidden_states=None, input_embeds=input_embeds, ) diff --git a/python/tokenspeed/runtime/layers/attention/backends/mha.py b/python/tokenspeed/runtime/layers/attention/backends/mha.py index 757e553ff..41c2b6ae8 100644 --- a/python/tokenspeed/runtime/layers/attention/backends/mha.py +++ b/python/tokenspeed/runtime/layers/attention/backends/mha.py @@ -119,6 +119,9 @@ def __init__(self, config: MHAConfig): return_lse=False, solution=self.kernel_solution, ) + # DFLASH draft: expand decode metadata to spec_num_tokens rows/request + # (whole block in one decode forward), with uniform non-causal seq_lens. + self.draft_block_decode = bool(getattr(config, "draft_block_decode", False)) # Forward metadata is initialized in the runner per forward call self.forward_decode_metadata: MHADecodeMetadata | None = None @@ -136,6 +139,9 @@ def init_forward_metadata( seq_lens: torch.Tensor, req_to_page: torch.Tensor, forward_mode: ForwardMode, + # Only consumed on the extend/mixed path; decode callers (e.g. the + # DFLASH draft and the cuda-graph wrapper's draft decode init) omit + # them, so they must be optional. extend_seq_lens: torch.Tensor | None = None, extend_seq_lens_cpu: torch.Tensor | None = None, extend_prefix_lens: torch.Tensor | None = None, @@ -198,10 +204,34 @@ def init_forward_metadata( seq_lens=seq_lens, ) else: - self.forward_decode_metadata = MHADecodeMetadata( - page_table=page_table, - seq_lens=seq_lens, - ) + if self.draft_block_decode and self.spec_num_tokens > 1: + # DFLASH drafts a whole block in one decode forward; the decode + # kernel keys masking off max_seqlen_q, so expand each request + # into spec_num_tokens rows with the SAME full seq_len. That + # makes max_seqlen_q == 1 per row, so every block query attends + # over the entire block (non-causal block-diffusion drafting). + # Target verify keeps the unexpanded multi-query decode path. + expanded_page_table, expanded_seq_lens = ( + self._make_spec_metadata_buffers( + bs, + page_table.device, + ) + ) + self._fill_spec_metadata_uniform( + expanded_page_table, + expanded_seq_lens, + page_table, + seq_lens, + ) + self.forward_decode_metadata = MHADecodeMetadata( + page_table=expanded_page_table, + seq_lens=expanded_seq_lens, + ) + else: + self.forward_decode_metadata = MHADecodeMetadata( + page_table=page_table, + seq_lens=seq_lens, + ) def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor): assert ( @@ -214,6 +244,18 @@ def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor): ) self.cuda_graph_decode_metadata = {} + if self.draft_block_decode and self.spec_num_tokens > 1: + # DFLASH draft block: expand to spec_num_tokens decode rows per + # request (one row per block position), so max_seqlen_q == 1 per row + # and every block query attends over the whole block (non-causal). + self.cuda_graph_page_table, self.cuda_graph_seq_lens = ( + self._make_spec_metadata_buffers(max_bs, self.device) + ) + self.cuda_graph_page_table.zero_() + # seq_lens are filled from the live draft length inside the captured + # graph; seed a valid baseline so any pre-broadcast read stays in range. + self.cuda_graph_seq_lens.fill_(self.spec_num_tokens) + return self.cuda_graph_page_table = torch.zeros( (max_bs, self.max_num_pages), dtype=torch.int32, device=self.device ) @@ -234,12 +276,24 @@ def init_forward_metadata_capture_cuda_graph( ): assert not forward_mode.is_extend_or_mixed() - metadata = MHADecodeMetadata( - page_table=self.cuda_graph_page_table[:bs, :], - seq_lens=self.cuda_graph_seq_lens[:bs], - ) - if self.spec_num_tokens > 1 and not self.is_draft: - metadata.seq_lens.copy_(seq_lens[:bs].clamp_min(self.spec_num_tokens)) + if self.draft_block_decode and self.spec_num_tokens > 1: + # DFLASH draft block: spec_num_tokens decode rows per request. + expanded_bs = bs * self.spec_num_tokens + metadata = MHADecodeMetadata( + page_table=self.cuda_graph_page_table[:expanded_bs, :], + seq_lens=self.cuda_graph_seq_lens[:expanded_bs], + ) + # Uniform non-causal seq_lens are written by the drafter inside the + # captured graph (see fill_block_decode_seq_lens); seed a safe + # baseline for the capture run before that op records. + metadata.seq_lens.fill_(self.spec_num_tokens) + else: + metadata = MHADecodeMetadata( + page_table=self.cuda_graph_page_table[:bs, :], + seq_lens=self.cuda_graph_seq_lens[:bs], + ) + if self.spec_num_tokens > 1 and not self.is_draft: + metadata.seq_lens.copy_(seq_lens[:bs].clamp_min(self.spec_num_tokens)) self.cuda_graph_decode_metadata[bs] = metadata self.forward_decode_metadata = metadata @@ -266,10 +320,39 @@ def init_forward_metadata_replay_cuda_graph( ) if self.spec_num_tokens > 1 and not self.is_draft: self.cuda_graph_seq_lens[:bs].copy_(seq_lens[:bs]) + elif self.draft_block_decode: + # DFLASH draft: replicate each request's page table to its + # spec_num_tokens block rows. The block-end seq_lens are filled by + # the drafter inside the captured graph, so they are not touched + # here (they re-derive from the live draft length on every replay). + base_page_table = req_to_page[req_pool_indices[:bs], : self.max_num_pages] + self.cuda_graph_page_table[: bs * self.spec_num_tokens, :].view( + bs, self.spec_num_tokens, self.max_num_pages + ).copy_(base_page_table[:, None, :]) if bs in self.cuda_graph_decode_metadata: self.forward_decode_metadata = self.cuda_graph_decode_metadata[bs] + def fill_block_decode_seq_lens(self, bs: int, block_seq_lens: torch.Tensor) -> None: + """DFLASH: broadcast each request's block-end length to its + spec_num_tokens cuda-graph decode rows (uniform, non-causal). + + Called by the drafter inside the captured graph so that on every replay + the expanded seq_lens re-derive from the live draft length (which is + recomputed in-graph from the target's accept lengths). + + Args: + bs: Number of draft requests. + block_seq_lens: ``[bs]`` per-request block-end lengths + (prefix + spec_num_tokens). + """ + spec = self.spec_num_tokens + self.cuda_graph_seq_lens[: bs * spec].view(bs, spec).copy_( + block_seq_lens[:bs] + .clamp(self.spec_num_tokens, self.max_context_len) + .unsqueeze(1) + ) + # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ @@ -524,6 +607,46 @@ def _get_kv_cache(self, layer: PagedAttention, token_to_kv_pool): ) return k_cache, v_cache + def _make_spec_metadata_buffers( + self, + bs: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + expanded_bs = bs * self.spec_num_tokens + cuda_graph_page_table = torch.empty( + (expanded_bs, self.max_num_pages), + dtype=torch.int32, + device=device, + ) + cuda_graph_seq_lens = torch.empty( + (expanded_bs,), + dtype=torch.int32, + device=device, + ) + return (cuda_graph_page_table, cuda_graph_seq_lens) + + def _fill_spec_metadata_uniform( + self, + expanded_page_table: torch.Tensor, + expanded_seq_lens: torch.Tensor, + page_table: torch.Tensor, + seq_lens: torch.Tensor, + ): + """Expand spec metadata with a uniform (non-causal) seq_len per row. + + Replicates the full seq_len to all spec_num_tokens rows of a request so + each row decodes with max_seqlen_q == 1 over the whole block. Used by the + DFLASH drafter so every block query attends over the entire block + (non-causal block-diffusion drafting), as opposed to the target's + unexpanded causal multi-query verify path. + """ + bs = seq_lens.shape[0] + spec_num_tokens = self.spec_num_tokens + expanded_page_table = expanded_page_table.view( + bs, spec_num_tokens, self.max_num_pages + ) + expanded_page_table.copy_(page_table[:, None, :]) + expanded_seq_lens.view(bs, spec_num_tokens).copy_(seq_lens[:, None]) for _backend_name in _KERNEL_SOLUTION_BY_BACKEND: register_backend(_backend_name, {AttentionArch.MHA}, MHAAttnBackend) diff --git a/python/tokenspeed/runtime/layers/attention/configs/base.py b/python/tokenspeed/runtime/layers/attention/configs/base.py index d6556045a..3bb2cb397 100644 --- a/python/tokenspeed/runtime/layers/attention/configs/base.py +++ b/python/tokenspeed/runtime/layers/attention/configs/base.py @@ -58,6 +58,10 @@ class BaseAttnConfig: speculative_num_steps: int = 0 speculative_num_draft_tokens: int = 1 is_draft: bool = False + # DFLASH drafts a whole block in one decode forward (q_len = spec_num_tokens + # per request) instead of Eagle/MTP's per-step single-token decode. Backends + # use this to expand decode metadata to spec_num_tokens rows per request. + draft_block_decode: bool = False @classmethod def generate( diff --git a/python/tokenspeed/runtime/layers/attention/configs/mha.py b/python/tokenspeed/runtime/layers/attention/configs/mha.py index 37fad45fc..1d9c5ca16 100644 --- a/python/tokenspeed/runtime/layers/attention/configs/mha.py +++ b/python/tokenspeed/runtime/layers/attention/configs/mha.py @@ -46,7 +46,10 @@ def generate( speculative_num_draft_tokens=server_args.speculative_num_draft_tokens, ) kv_cache_dtype = server_args.kv_cache_dtype - if is_draft and server_args.speculative_algorithm == "DFLASH": + draft_block_decode = bool( + is_draft and server_args.speculative_algorithm == "DFLASH" + ) + if draft_block_decode: kv_cache_dtype = "bfloat16" return cls( @@ -69,6 +72,7 @@ def generate( max_graph_bs=server_args.max_cudagraph_capture_size, kv_cache_quant_method=server_args.kv_cache_quant_method, is_draft=is_draft, + draft_block_decode=draft_block_decode, **kwargs, ) diff --git a/python/tokenspeed/runtime/models/dflash.py b/python/tokenspeed/runtime/models/dflash.py index 993dfc970..8c19fb2d3 100644 --- a/python/tokenspeed/runtime/models/dflash.py +++ b/python/tokenspeed/runtime/models/dflash.py @@ -276,9 +276,7 @@ def forward( elif ( ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"] ): - hidden_states = all_reduce( - hidden_states, self.mapping.dense.tp_rank, self.mapping.dense.tp_group - ) + hidden_states = all_reduce(hidden_states, self.mapping.dense.tp_group) hidden_states, residual = self.input_layernorm(hidden_states, residual) else: hidden_states, residual, *_ = ( @@ -298,9 +296,7 @@ def forward( ) if ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"]: - hidden_states = all_reduce( - hidden_states, self.mapping.attn.tp_rank, self.mapping.attn.tp_group - ) + hidden_states = all_reduce(hidden_states, self.mapping.attn.tp_group) hidden_states, residual = self.post_attention_layernorm( hidden_states, residual ) @@ -364,7 +360,7 @@ def forward( input_ids: torch.Tensor, positions: torch.Tensor, out_cache_loc: torch.Tensor, - input_lengths: torch.Tensor, + input_lengths: torch.Tensor | None = None, input_embeds: torch.Tensor | None = None, **kwargs, ) -> LogitsProcessorOutput: From ed5e0fc1c1a32a1b06bfda7551c8435a83d21c73 Mon Sep 17 00:00:00 2001 From: Yue Weng <25103990+yweng0828@users.noreply.github.com> Date: Thu, 18 Jun 2026 01:28:10 -0700 Subject: [PATCH 3/4] add ut --- test/runtime/models/test_kimi_models.py | 44 +++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/test/runtime/models/test_kimi_models.py b/test/runtime/models/test_kimi_models.py index c535e1f0f..841070004 100644 --- a/test/runtime/models/test_kimi_models.py +++ b/test/runtime/models/test_kimi_models.py @@ -10,6 +10,7 @@ python3 -m unittest models.test_kimi_models.TestKimiK25.test_base -v python3 -m unittest models.test_kimi_models.TestKimiK25.test_tsckpt_eagle3 -v python3 -m unittest models.test_kimi_models.TestKimiK25.test_nvckpt_eagle3 -v + python3 -m unittest models.test_kimi_models.TestKimiK25.test_dflash -v Environment (all optional): KIMI_K25_MODEL HF model id or path (default: nvidia/Kimi-K2.5-NVFP4) @@ -17,6 +18,7 @@ KIMI_K25_WORLD_SIZE GPU count (default: 4) KIMI_K25_DRAFT_MODEL EAGLE3 draft repo (default: lightseekorg/kimi-k2.5-eagle3) KIMI_K25_MLA_DRAFT_MODEL MLA EAGLE3 draft repo (default: nvidia/Kimi-K2.5-Thinking-Eagle3) + KIMI_K25_DFLASH_DRAFT_MODEL Native DFLASH draft repo (default: z-lab/Kimi-K2.5-DFlash) """ import dataclasses @@ -37,6 +39,9 @@ MLA_DRAFT_MODEL = os.environ.get( "KIMI_K25_MLA_DRAFT_MODEL", "nvidia/Kimi-K2.5-Thinking-Eagle3" ) +DFLASH_DRAFT_MODEL = os.environ.get( + "KIMI_K25_DFLASH_DRAFT_MODEL", "z-lab/Kimi-K2.5-DFlash" +) TIMEOUT = 600 _server_port = 22000 @@ -207,6 +212,41 @@ class MeshCase: "trtllm_mla", ), ), + "dflash": MeshCase( + "dflash", + ( + # Native DFLASH drafts a whole block per decode step. The main model + # runs MLA (tokenspeed_mla), but the z-lab/Kimi-K2.5-DFlash draft is + # a Qwen3 MHA/GQA model, so the drafter must use an MHA-family + # backend (fa4); tokenspeed_mla is MLA-only and would reject it. + "--attention-backend", + "tokenspeed_mla", + "--moe-backend", + "flashinfer_trtllm", + "--kv-cache-dtype", + "fp8_e4m3", + "--max-prefill-tokens", + "8192", + "--chunked-prefill-size", + "8192", + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + DFLASH_DRAFT_MODEL, + "--speculative-num-steps", + "7", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "8", + "--speculative-draft-model-quantization", + "unquant", + "--drafter-attention-backend", + "fa4", + "--sampling-backend", + "greedy", + ), + ), } @@ -246,6 +286,10 @@ def test_nvckpt_eagle3(self): """Kimi K2.5 with MLA EAGLE3 draft (trtllm_mla drafter + FP8 KV cache).""" self._run_quality_checks(MESH_CASES["nvckpt_eagle3"]) + def test_dflash(self): + """Kimi K2.5 with native DFLASH block-diffusion draft (fa4 drafter).""" + self._run_quality_checks(MESH_CASES["dflash"]) + if __name__ == "__main__": unittest.main() From aa31274cdc13b187bf2cd2743ef2c0c26ee9f9c2 Mon Sep 17 00:00:00 2001 From: Yue Weng <25103990+yweng0828@users.noreply.github.com> Date: Tue, 23 Jun 2026 05:19:31 -0700 Subject: [PATCH 4/4] fix ima --- .../runtime/execution/drafter/dflash.py | 79 +++++++++++++++---- .../runtime/layers/attention/backends/mha.py | 13 ++- 2 files changed, 74 insertions(+), 18 deletions(-) diff --git a/python/tokenspeed/runtime/execution/drafter/dflash.py b/python/tokenspeed/runtime/execution/drafter/dflash.py index fb2c54e06..8a8498eb1 100644 --- a/python/tokenspeed/runtime/execution/drafter/dflash.py +++ b/python/tokenspeed/runtime/execution/drafter/dflash.py @@ -152,6 +152,57 @@ def bind_target_model(self, target_model) -> None: self.lm_head = target_model.lm_head self.logits_processor = language_model.logits_processor + def _greedy_gather_capacity(self) -> int: + """Max element count for the greedy head's tensor-parallel all-gather + scratch: a full ``max_bs`` decode block. + + The greedy head samples the last ``spec_num_tokens - 1`` block + positions per request and all-gathers them across the TP group, so the + worst case is ``tp_size * max_bs * (spec_num_tokens - 1)``. + """ + tp_size = int(self.logits_processor.tp_size) + return tp_size * self.input_buffers.max_bs * max(self.spec_num_tokens - 1, 1) + + def _ensure_greedy_gather_buffers( + self, + max_dtype: torch.dtype, + ids_dtype: torch.dtype, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Lazily create the greedy all-gather scratch ONCE at its maximum + capacity, then reuse it in place for every batch size. + + Sizing to the max ``max_bs`` block (rather than growing per batch size) + is required for CUDA-graph correctness. Graphs are captured for + increasing batch sizes (``[1, 2, ..., max_bs]``); a buffer grown lazily + would be freed and reallocated when a larger bs needs more room, leaving + every smaller-bs graph captured earlier with an + ``all_gather_into_tensor`` recorded against freed memory. On replay + those small-bs decode steps read garbage (out-of-vocab) draft token ids, + which flow into the next verify forward's embedding lookup and trigger a + CUDA illegal memory access. A fixed max-capacity buffer is allocated + during warmup (before capture) and shared by every captured graph. + + Returns the (max, id) scratch tensors; callers slice ``[:needed]``. + """ + cap = self._greedy_gather_capacity() + if ( + self._greedy_gathered_max is None + or self._greedy_gathered_ids is None + or self._greedy_gather_cap < cap + or self._greedy_gathered_max.dtype != max_dtype + or self._greedy_gathered_max.device != device + or self._greedy_gathered_ids.dtype != ids_dtype + ): + self._greedy_gathered_max = torch.empty( + (cap,), dtype=max_dtype, device=device + ) + self._greedy_gathered_ids = torch.empty( + (cap,), dtype=ids_dtype, device=device + ) + self._greedy_gather_cap = cap + return self._greedy_gathered_max, self._greedy_gathered_ids + def _greedy_sample_from_vocab_parallel_head( self, hidden_states: torch.Tensor, @@ -221,23 +272,11 @@ def _greedy_sample_from_vocab_parallel_head( return global_ids.to(torch.int32) needed = tp_size * chunk_len - if ( - self._greedy_gather_cap < needed - or self._greedy_gathered_max is None - or self._greedy_gathered_ids is None - or self._greedy_gathered_max.dtype != local_max.dtype - or self._greedy_gathered_max.device != hidden_states.device - ): - self._greedy_gathered_max = torch.empty( - (needed,), dtype=local_max.dtype, device=hidden_states.device - ) - self._greedy_gathered_ids = torch.empty( - (needed,), dtype=global_ids.dtype, device=hidden_states.device - ) - self._greedy_gather_cap = needed - - gathered_max = self._greedy_gathered_max[:needed] - gathered_ids = self._greedy_gathered_ids[:needed] + gathered_max, gathered_ids = self._ensure_greedy_gather_buffers( + local_max.dtype, global_ids.dtype, hidden_states.device + ) + gathered_max = gathered_max[:needed] + gathered_ids = gathered_ids[:needed] all_gather_into_tensor( gathered_max, local_max.contiguous(), @@ -497,6 +536,12 @@ def _draft_native(self, current_tokens: torch.Tensor) -> torch.Tensor: draft_hidden[:, 1:, :].reshape(-1, self.hidden_size) ) next_tokens[:, 1:] = sampled.view(bs, self.spec_num_tokens - 1) + # Defense-in-depth: keep draft ids non-negative before they are written + # to future_input_map and embedded by the next verify forward, mirroring + # the EAGLE drafter's draft_ids.clamp_(min=0) guard. A negative id (the + # -1 NaN sentinel) would otherwise index the embedding table out of + # bounds (CUDA illegal memory access). + next_tokens.clamp_(min=0) return next_tokens @nvtx_range("drafter:dflash", color="purple") diff --git a/python/tokenspeed/runtime/layers/attention/backends/mha.py b/python/tokenspeed/runtime/layers/attention/backends/mha.py index 41c2b6ae8..52cd6fd68 100644 --- a/python/tokenspeed/runtime/layers/attention/backends/mha.py +++ b/python/tokenspeed/runtime/layers/attention/backends/mha.py @@ -646,7 +646,18 @@ def _fill_spec_metadata_uniform( bs, spec_num_tokens, self.max_num_pages ) expanded_page_table.copy_(page_table[:, None, :]) - expanded_seq_lens.view(bs, spec_num_tokens).copy_(seq_lens[:, None]) + # Clamp to max_context_len so the draft decode never asks the attention + # kernel for more than max_num_pages worth of page-table columns. The + # block-end length is prefix + spec_num_tokens, which can exceed + # max_context_len for a request near the context limit; without the + # clamp the kernel reads page_table[:, >= max_num_pages] out of bounds + # (CUDA illegal memory access). Mirrors fill_block_decode_seq_lens on the + # cuda-graph path (this eager path is taken by mixed prefill+decode + # batches even when cuda graphs are enabled). + expanded_seq_lens.view(bs, spec_num_tokens).copy_( + seq_lens.clamp(spec_num_tokens, self.max_context_len)[:, None] + ) + for _backend_name in _KERNEL_SOLUTION_BY_BACKEND: register_backend(_backend_name, {AttentionArch.MHA}, MHAAttnBackend)