artifacts/dft01-compact-draft-v0/RESULTS.json

{
  "task": "DFT-01",
  "timestamp": "2026-06-11T22:39:36.946981Z",
  "drafter_architecture": {
    "config_model_type": "gemma4_assistant (Gemma4AssistantForCausalLM)",
    "drafter_own_lm_head": false,
    "tie_word_embeddings": true,
    "lm_head_weight_source": "model.embed_tokens.weight (tied)",
    "lm_head_vocab_size": 262144,
    "lm_head_shape": [
      262144,
      256
    ],
    "masked_embedding": {
      "num_centroids": 2048,
      "vocab_size_per_centroid": 128,
      "centroid_intermediate_top_k": 64,
      "num_selected": 8192
    },
    "drafter_hidden_size": 256,
    "backbone_hidden_size": 2560,
    "num_layers": 4,
    "speculative_config": "{method:mtp,model:/tmp/qat-assistant,num_speculative_tokens:7}",
    "checkpoint_size_bytes": 159138240,
    "checkpoint_sha256": "ed159e334999fd6b5f2d0dbad026346d4efac89eb7c6f55c5cdb042eca5dd18e",
    "note": "No lm_head tensor in safetensors; tied entirely through model.embed_tokens.weight [262144,256]"
  },
  "scatter_fill_value": {
    "pck04_fill": "float(-inf) \u2014 torch.full(..., float(-inf))",
    "drafter_masked_embedding_fill": "torch.finfo(hidden_states.dtype).min (used in forward(), NOT in get_top_tokens())",
    "note": "PCK04 scatter patch only wraps Gemma4ForCausalLM.compute_logits (target); drafter uses Gemma4AssistantForCausalLM which is NOT patched"
  },
  "scatter_fires_on_draft_steps": false,
  "scatter_fire_analysis": {
    "target_path": "GPUModelRunner.execute_model -> Gemma4ForCausalLM.compute_logits [patched by pck04] -> scatter 131072->262144 full vocab",
    "draft_path": "Gemma4Proposer._greedy_sample -> (CUDA-graphed) Gemma4MTPMaskedEmbedder.get_top_tokens -> sparse argmax over [1, 8192] -> returns vocab_id directly, NO 262k allocation",
    "pck04_patch_scope": "Only Gemma4ForCausalLM (target); Gemma4AssistantForCausalLM (drafter) is not patched",
    "draft_vocab_is_full_262144": true,
    "draft_uses_pruned_keepset": false,
    "note": "The drafter already does compact argmax natively. The vocabulary it argmaxes over is the FULL 262144 vocab embed_tokens weight (not the 131072 keepset). It selects from 8192 centroid-filtered candidates via sparse dot-product, then returns the winning token ID directly without scatter."
  },
  "hypothesis_verdict": "DEAD \u2014 the optimization already exists in the current stack. The drafter NEVER materializes a 262144-wide logits tensor during decode. The Gemma4Proposer._greedy_sample calls model.get_top_tokens() (centroids CUDA-graphed for size=1), which performs a 2-phase sparse argmax: centroid scoring [1,2048] -> topk 64 centroids -> gather 8192 token embeddings -> dot product -> argmax -> return token_id. No 262144 scatter, no allocation, no PCK04 involvement.",
  "compact_draft_patch_needed": false,
  "compact_draft_already_implemented": true,
  "implementation_location": {
    "drafter_model": "vllm/model_executor/models/gemma4_mtp.py:Gemma4MTPMaskedEmbedder.get_top_tokens()",
    "proposer_dispatch": "vllm/v1/spec_decode/gemma4.py:Gemma4Proposer._greedy_sample()",
    "cuda_graph_capture": "vllm/v1/spec_decode/gemma4.py:Gemma4Proposer._setup_centroids_cuda_graphs()",
    "sitecustomize_fused_argmax": "sitecustomize.py:_apply_fused_top_token_patch() patches Gemma4MTPMaskedEmbedder.get_top_tokens to use Triton fused sparse argmax kernel"
  },
  "sitecustomize_fused_argmax_note": "FUSED_SPARSE_ARGMAX=1 in manifest patches get_top_tokens with a Triton 2-pass (blocks+reduce) kernel. FUSED_SPARSE_ARGMAX_REQUIRE=1 makes it fail-loud. This IS already the compact drafter path.",
  "potential_remaining_opportunity": {
    "description": "The drafter lm_head_weight is the FULL 262144-row embed_tokens. The keepset covers 131072 tokens (~99.41% of corpus). If we were to restrict the drafter vocabulary to the 131072 keepset, the embedding gather (8192 token embeddings of dim 256) would be cheaper to fetch. But: (a) the drafter is CUDA-graphed at size=1 and centroids CUDA graphs are captured once, so the lm_head_weight access is static; (b) the drafter proposes from the full 262144 vocab \u2014 restricting it to the keepset would prevent proposals of non-keepset tokens, which the target would reject anyway (they get -inf logprob from pck04 scatter), potentially improving acceptance rate or being a no-op; (c) this requires rewriting the drafter centroid/token_ordering structures.",
    "verdict": "Out of scope for DFT-01; would require drafter checkpoint surgery, not a serving-side patch"
  },
  "blockers": null,
  "artifacts_produced": [
    "/workspace/dft01/RESULTS.json"
  ]
}