Feat/flash attn attn bias - #68
Merged
Merged
Conversation
) #1 — Draft KV cache. block_diffusion_attention gains return_ctx_kv to expose this step's projected+roped context KV; DraftKVCache accumulates it per layer. dflash_decoder_layer_cached / dflash_draft_forward_cached thread the cache so the draft attends to the full accumulated context instead of recomputing it. Verified: cached(accumulated) == non-cached(full context) to 1e-3, and the cache accumulates/advances correctly across steps. #2 — Non-greedy sampling + distribution-preserving acceptance. make_sampler (temperature / top-k / top-p, rng-reproducible), sampler_probs (matching truncated distribution), and dflash_speculative_verify (Leviathan rule: accept d_i w.p. min(1, p_t/p_d); on reject draw from normalize(relu(p_t-p_d)); bonus from the target's next-position distribution). Verified: greedy == argmax, top-k restricts support, draft==target accepts all, and the speculative-sampling theorem — the emitted token's marginal equals the target distribution (Monte Carlo, 40k draws, max abs err < 0.02). tests/unit/test_dflash_cached_sampling.py (7). ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…odel (#4) #4 — tessera.dflash_reference.ReferenceDecoderLM: a small numpy causal decoder (pre-norm MHA + SwiGLU, rope, tied/untied LM head) with a multi-layer hidden tap (the DFlash conditioning signal) and a stateless forward() that is the greedy-AR ground truth. random_decoder_lm builds one with small random weights. #3 — stateful KV cache + rollback: step(tokens) does causal cached decoding and appends roped-K/V per layer; rollback(n) drops the over-speculated tail. Verified that incremental step() (in 3 chunks) reproduces the stateless full-sequence forward to 1e-3, and that rollback restores exact cache state. dflash_generate_cached ties it together: cached draft (#1) + stateful target with rollback (#3) + greedy or rejection sampling (#2). Verified the whole efficient loop reproduces greedy AR exactly, is block-size independent, and sampling is reproducible + in-vocab. tests/unit/test_dflash_reference_target.py (5). ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…oint I/O (#7) #5 — attention_fn threaded through dflash_decoder_layer / dflash_draft_forward (+ cached variants) so the whole draft forward runs its attention on the Apple GPU metal_runtime lane via apple_gpu_attention_fn. Verified the whole draft (2 layers) matches the numpy reference on Metal (rtol/atol 1e-3). The matmul- heavy projections/MLP/LM-head stay host-side (GPU gather/embedding is the remaining blocker for a single fully-jitted artifact). #9a — position-weighted block training loss: dflash_position_weights (wₖ = exp(-k/γ), normalized), dflash_block_loss (mean/sum/none) and the explicit gradient dflash_block_loss_grad. Verified the gradient vs finite differences (<1e-7), that a grad step lowers the loss, and reduction consistency. #7 — checkpoint I/O (tessera.dflash_io): a dependency-free safetensors reader/writer + HF state-dict <-> DFlashWeights mapping (transposing the nn.Linear (out,in) weights to the x@W (in,out) convention; embedding/LM head supplied from the target). load_dflash_weights reads a z-lab/*-DFlash draft; verified safetensors round-trip, the (out,in) transpose, and that round-tripped weights produce identical draft logits. tests: test_dflash_train_io.py (7) + #5 GPU draft case. ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…uler (#9c/#9d) #6 — DFlashDraft(nn.Module): holds every draft tensor as a Parameter (so it participates in parameters()/state_dict/.to(dtype)), forwards through the functional draft (cached or not), from_weights()/to_weights() round-trip. Verified module forward == functional (<1e-5), 5 + 11*N params registered, weight round-trip. #9b — RotatingDraftKVCache: bounds the draft context cache to the last max_size tokens (the draft analogue of MLX RotatingKVCache for sliding layers). Verified it caps per-layer length and, when unbounded, is identical to DraftKVCache. #9c/#9d — tessera.dflash_serve: dflash_generate_text (string-in/out via any encode/decode tokenizer) and DFlashScheduler (holds draft + stateful target, serves generation requests, greedy == AR). Verified scheduler greedy == AR and generate_text round-trips through a tokenizer. tests/unit/test_dflash_module_serve.py (5). ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
#8 — annotate block_diffusion_attention's GQA: repeat is numerically exact; the native flash_attn_gqa kernel doesn't support DFlash's concat-context+proposal KV with an additive bias, so the reference materializes the repeat (no code change — correctness is unaffected). MASTER_AUDIT records DFlash integration items 1–9 as landed, with the two remaining gates flagged as external (real-checkpoint numerical parity needs a network download; a single fully-jitted GPU draft artifact needs GPU gather). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- README status table: new "Speculative decoding — attn_bias substrate + DFlash block-diffusion draft" row (honest status: Python reference + attention core on Apple GPU metal_runtime; greedy spec-decode == greedy AR proven vs the MLX reference; real-checkpoint parity + fully-jitted GPU draft are external gates). - README: refresh stale Apple C ABI counts to the generated truth (256→264 symbols, 109→112 kernel families). - New docs/dflash.md: user-facing overview — the attn_bias substrate, the module map (dflash / dflash_reference / dflash_io / dflash_serve), a quick start, what's proven (per-test), and the external gates. Linked from the README doc index. docs lint passes; all links resolve. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…CANONICAL_API PYTHON_API_SPEC.md: - flash_attn signature + parameter table gain attn_bias (additive (B,Sq,Sk) mask, Apple GPU flash_attn_bias_* / metal_runtime, causal+bias, broadcast fallback, positional-bias VJP). - Module hierarchy lists tessera.dflash / dflash_reference / dflash_io / dflash_serve. - New §18 "Speculative Decoding (DFlash)" documents the full public surface across the four modules + nn.functional.block_diffusion_attention / mask_token_block; TOC + Appendix A symbol index updated. CANONICAL_API.md: - flash_attn ops row gains attn_bias; functional table gains block_diffusion_attention + mask_token_block; new "tessera.dflash — Speculative Decoding (DFlash)" section with canonical names (one per concept) + a quick-start. Verified: check_spec_sync + docs lint pass; every documented symbol exists and every module __all__ symbol is documented (zero drift, confirmed programmatically). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
A programmatic sweep of the public surface vs both API docs found the surface ~99% documented with concentrated gaps; this closes them to zero. tessera.ops (312 ops): added the 2 missing — bmm (batched matmul + broadcast, Apple GPU metal_runtime) and fake_quantize (QAT STE) — to both ops tables. tessera.nn (77 public attrs): added the 10 missing functional layers (linear_general, lora_linear, spectral_norm, conv_transpose, avg/max/min/adaptive pool, gru_cell, simple_rnn_cell, bidirectional_scan) to CANONICAL's functional table, and the 10 missing Module classes (LinearGeneral, Einsum, LoRALinear, ConvTranspose1d/ConvTranspose, SpectralNorm, GRUCell/SimpleRNNCell, NativeSparseAttention, MixtureOfRecursions) to the stateful class table, with accurate constructor/forward signatures. Verified programmatically: tessera.ops, tessera.nn, and nn.functional.__all__ now have ZERO undocumented public symbols. check_spec_sync + docs lint pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8d2a744b24
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| def Tessera_FlashAttnOp : Op< | ||
| Tessera_Dialect, "flash_attn", | ||
| [Pure] | ||
| [Pure, AttrSizedOperandSegments] |
There was a problem hiding this comment.
With this change tessera.flash_attn now has the AttrSizedOperandSegments trait, but the Python Graph IR printer still emits ordinary 3-operand flash-attention ops without an operandSegmentSizes attribute (I checked python/tessera/compiler/graph_ir.py::IROp.to_mlir, which only synthesizes that attr for tessera.rl.ppo_policy_loss). That means existing Python-generated bias-free tessera.flash_attn %Q, %K, %V ... MLIR will fail the trait's required segment-size verification/parsing unless callers hand-add <{operandSegmentSizes = array<i32: 1, 1, 1, 0>}>, as the updated tests now do manually. Please add segment-size emission/inference for flash_attn when introducing this trait.
Useful? React with 👍 / 👎.
Comment on lines
+784
to
+785
| first = int(np.asarray(sample(logits[:, -1:])).reshape(-1)[0]) | ||
| tokens.append(first) |
There was a problem hiding this comment.
When max_new_tokens is 0, this unconditional prefill still samples and appends first, so dflash_generate_cached(..., max_new_tokens=0) returns one generated token instead of the prompt unchanged; the stateless dflash_generate has the same pattern. This also advances the target cache for a request that should be a no-op, which can surprise serving code that uses zero-length continuations for warmup or accounting.
Useful? React with 👍 / 👎.
gstoner
force-pushed
the
feat/flash-attn-attn-bias
branch
from
June 13, 2026 12:44
8d2a744 to
9628c7b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.