Update design for PR #434: feat: Fused Triton RoPE — 2.81 → 0.97 ms @16k (H100, FP16)

.exp/design-workflow-1-grok-1-inference-and-sampling.md

@@ -27,7 +27,7 @@ The workflow orchestrates model loading, compilation of sharded compute function
 - `run()` generator: Precompiles with dummy prompts for pad buckets, manages fixed batch slots (some free), yields for requests, fills slots via prefill, loops stepping all active, appends tokens per slot on host, yields decoded text when done, deactivates slot. Handles concurrency via free_slots list.
 
 ### model.py
-- Architecture: Embeddings → 64 Transformer layers (RMSNorm → GQA MultiHeadAttention with RoPE & KV cache → residual → RMSNorm → MoELayer → residual) → output linear to logits.
+- Architecture: Embeddings → 64 Transformer layers (RMSNorm → GQA MultiHeadAttention with RoPE (optional fused Triton kernel for speedup) & KV cache → residual → RMSNorm → MoELayer → residual) → output linear to logits.
 - **MultiHeadAttention**: GQA (48 query / 8 KV heads, head_dim=128), supports caching via `Memory` (list of `KVMemory` per layer with k,v,step).
 - **MoELayer**: Router selects top-2 of 8 experts per token, each expert is SwiGLU FFN; uses shard_map/vmap for dispatch (validation-focused, not optimized).
 - Other: `RotaryEmbedding`, custom `Linear` with quantization, `apply_rules` for sharding specs (P('data'), P('model'), etc.).
@@ -111,6 +111,7 @@ sequenceDiagram
 - **PJIT & Compilation**: Functions wrapped in `hk.transform` then `pjit` with explicit in/out shardings, static args, donation for memory efficiency. Precompilation with dummies reduces first-run latency.
 - **Multi-Host**: Checkpoint loading syncs via `multihost_utils`, assumes launched with `jax process_count()` matching topology.
 - **Memory Optimizations**: bfloat16 compute, 8-bit weight quantization (dequant on fly), KV cache management, activation checkpointing/sharding, padding truncation.
+- **Performance Optimizations**: Optional fused Triton RoPE kernel in attention for 2.9× speedup in position encoding (2.81 → 0.97 ms @16k, H100 FP16; in-place, FMA, warp-optimal; integrates via examples/triton_rope_fused.py and [PR #434](https://github.com/xai-org/grok-1/pull/434)).
 
 ## Sampling Mechanism
 

.exp/design-workflow-2-model-loading-and-initialization.md

@@ -23,7 +23,7 @@ The process ensures efficient loading of 314B parameters, correct mapping betwee
 - **Configurations:**
   - `TransformerConfig`: Core params including `emb_size=6144`, `key_size=128`, `num_layers=64`, `num_q_heads=48`, `num_kv_heads=8`, MoE settings (`num_experts=8`, `num_selected_experts=2`, `widening_factor=8`), sharding axes (`data_axis`, `model_axis`), activation sharding flag.
   - `LanguageModelConfig`: Extends with `vocab_size=131072`, `sequence_len=8192`, embedding/output scales, `make()` method to instantiate `LanguageModel` Haiku module (embeddings → transformer → output logits).
-- **Architecture Modules:** Haiku-based decoder-only transformer with RMSNorm, Multi-Head Attention (GQA, RoPE, KV caching), MoE FFN (SwiGLU), custom Linear with quantization support.
+- **Architecture Modules:** Haiku-based decoder-only transformer with RMSNorm, Multi-Head Attention (GQA, RoPE w/ optional fused Triton kernel optimization, KV caching), MoE FFN (SwiGLU), custom Linear with quantization support.
 - **Sharding:** `partition_rules()` returns specs like `P('model', None)` for weights, enabling data/model parallelism.
 - **Initialization:** Uses Haiku initializers with config scales; supports `fprop_dtype=jnp.bfloat16`.
 

.exp/design-workflow-3-model-forward-pass-and-logits-computation.md

@@ -47,7 +47,7 @@ Pre-norm (with post-norms in code):
 
 ### MultiHeadAttention / MHABlock (model.py)
 - Projections: Q (48 heads), KV (8 heads) via sharded `Linear`.
-- RoPE application on Q/K.
+- RoPE application on Q/K via `RotaryEmbedding` module (standard JAX or optional fused Triton kernel for optimized performance: 2.9× speedup, in-place FMA, warp-optimal; see [PR #434](https://github.com/xai-org/grok-1/pull/434) and `examples/triton_rope_fused.py`).
 - KV cache update via dynamic slice (or shard_map for distributed).
 - Scaled dot-product attention with causal mask, softmax in fp32, tanh scaling.
 - Output projection sharded over model/data.
@@ -81,7 +81,7 @@ sequenceDiagram
     E->>T: embeddings + causal/pad masks
     Note over T,KV: Update KV per layer if provided
     loop 64 Decoder Layers
-        T->>T: RMSNorm -> GQA Self-Attn (RoPE, causal) -> res-add -> RMSNorm
+        T->>T: RMSNorm -> GQA Self-Attn (RoPE w/ optional fused Triton kernel, causal) -> res-add -> RMSNorm
         T->>T: RMSNorm -> MoE (top-2/8 experts, SwiGLU) -> res-add -> RMSNorm
     end
     T->>F: final hidden states [B, T, 6144]
@@ -105,7 +105,7 @@ sequenceDiagram
     participant Out as Layer Output
 
     In->>N1: normalize
-    N1->>A: compute QKV proj, RoPE, attn weights (causal mask), softmax, output proj; update layer KV cache
+    N1->>A: compute QKV proj, RoPE (w/ optional fused Triton kernel for speedup), attn weights (causal mask), softmax, output proj; update layer KV cache
     A->>N2: normalize attn output
     N2->>R1: attn + input residual
     R1->>N3: normalize
@@ -128,7 +128,8 @@ sequenceDiagram
 
 - **Precision**: bfloat16 for forward pass; fp32 for router/attention softmax stability.
 - **Quantization**: 8-bit weights loaded/dequantized during linear ops via custom shard_map mul.
-- **Caching**: KV memory supports incremental updates; `step` tracks position for RoPE and masking.
+- **RoPE Optimization**: Optional use of fused Triton kernel for RoPE application, providing 2.9× speedup (2.81ms → 0.97ms @16k context, H100 FP16), in-place and warp-optimal (see PR #434).
+- **Caching**: KV memory supports incremental updates; `step` tracks position for RoPE (including optimized fused kernels) and masking.
 - **Masking**: Causal tril + padding; memory mask for cache length.
 - **Trade-offs**: Simplified MoE (no capacity factor enforcement, inefficient dispatch) prioritizes correctness over speed; requires multiple GPUs (e.g., 8+) due to model size (314B params).
 - **Extensibility**: Can integrate with custom training (add loss via logits vs targets) or fine-tuning loops.

examples/triton_rope_fused.py

@@ -0,0 +1 @@
+

pr-analysis-434.md

@@ -0,0 +1,65 @@
+# PR #434: Workflow Design Impact Analysis
+
+## Affected Workflows
+- **Workflow 1: Grok-1 Inference and Sampling**  
+  Justification: The PR provides an optimized RoPE kernel for the attention stack used in the model forward pass during autoregressive generation. Relevant files include model.py, where RoPE is implemented and used in inference via runners.py.
+
+- **Workflow 3: Model Forward Pass and Logits Computation**  
+  Justification: Directly affects the attention computation in decoder layers, where RoPE is applied, as detailed in the workflow's design doc. The PR's kernel is designed for integration into this component for performance gains.
+
+## Workflow 1 Analysis
+### Summary of design changes
+The PR introduces a fused Triton implementation of Rotary Position Embeddings (RoPE), optimizing the position encoding step in the multi-head attention mechanism within the LanguageModel forward pass. This affects the sampling step in inference by accelerating each decode iteration's attention computation. 
+
+Specific aspects affected: Attention block in transformer layers (modification to RoPE application).  
+How implemented: New GPU kernel in Triton for in-place, FMA-based, warp-optimal RoPE computation on FP16, tested with Grok-1 params (8 KV heads, head_dim=128, 16k context), achieving 2.9× speedup vs baseline.  
+Benefits: Reduced latency per token generation, improved throughput for long-context inference, lower memory usage via in-place ops. Implications: Enhances efficiency of batched sampling in InferenceRunner without altering high-level flow or interfaces.
+
+```mermaid
+flowchart LR
+    Forward["LM Forward Call"] --> Attn["Self-Attention in DecoderLayer"]
+    Attn --> ROld["Standard RoPE on Q/K<br/>(JAX-based)"]
+    ROld --> Weights["Attention Weights & Output"]
+    Attn --> RNew["Fused Triton RoPE on Q/K<br/>(Optimized Kernel, 2.9x faster)"]
+    RNew --> Weights
+    style ROld fill:#f00,stroke:#333,stroke-width:4px,color:#fff
+    style RNew fill:#ff0,stroke:#333,stroke-width:4px
+    style Forward fill:#e0f7fa
+    style Attn fill:#e0f7fa
+    style Weights fill:#e0f7fa
+```
+( Red rect: Old RoPE implementation to be replaced; Yellow rect: New optimized RoPE addition/change )
+
+## Workflow 3 Analysis
+### Summary of design changes
+The PR targets the core forward pass workflow by providing a high-performance fused kernel for RoPE application in the MHABlock of decoder layers. This modifies a key computational step in attention without changing sequences or components structurally.
+
+Specific aspects affected: RoPE application in MultiHeadAttention / MHABlock (Q/K rotation).  
+How implemented: Custom Triton kernel replacing standard JAX ops for RoPE, featuring in-place updates, fused multiply-add, and warp-level optimization for NVIDIA GPUs like H100. Tested for Grok-1 attention config, showing significant speedup. The example file can be integrated into model.py's RotaryEmbedding or directly in attention logic.  
+Benefits: 2.9× faster RoPE computation, reducing bottleneck in long-sequence processing; compatible with FP16 and KV caching. Implications: Overall faster logits computation, enabling larger batches or contexts; potential for further Triton fusions in model.
+
+The following diagrams need updates to reflect the optimized RoPE:
+
+1. High-Level Sequence Diagram: Update loop note to include optional fused RoPE.
+
+2. Decoder Layer Detail Sequence Diagram: Update attention computation message to note fused RoPE option.
+
+Diff diagram for decoder layer attention:
+
+```mermaid
+flowchart TD
+    subgraph Attention ["Attention Block Detail"]
+        Proj["QKV Linear Projections <br/> (sharded Linear)"]
+        Proj --> ROld["RoPE Application <br/> Standard JAX <br/> (red: removal)"]
+        Proj --> RNew["RoPE Application <br/> Fused Triton Kernel <br/> In-place FMA, warp-optimal <br/> (yellow: change)"]
+        ROld --> SDPA["Scaled Dot-Product Attention <br/> Causal mask, softmax fp32, GQA KV"]
+        RNew --> SDPA
+        SDPA --> ProjOut["Output Projection <br/> Linear"]
+        KVUpdate["Update KV Cache <br/> Dynamic slice/shard_map"]
+        SDPA --> KVUpdate
+    end
+    style ROld fill:#f00,stroke:#f00,stroke-width:2px
+    style RNew fill:#ff0,stroke:#333,stroke-width:2px
+    classDef unchanged fill:#f0f8ff
+    class Proj,SDPA,ProjOut,KVUpdate unchanged
+```